From f46b4a6fa263d7cf51bc8f3ceb2a69d2c1e83fdd Mon Sep 17 00:00:00 2001 From: pufferffish Date: Fri, 14 Jun 2024 19:56:35 +0100 Subject: [PATCH 01/29] implement the vulkan C backend --- gpu/gpu_vulkan.c | 126 +++++++++++++++++++++++++++++++++++++++++++++++ gpu/gpu_vulkan.h | 17 +++++++ 2 files changed, 143 insertions(+) create mode 100644 gpu/gpu_vulkan.c create mode 100644 gpu/gpu_vulkan.h diff --git a/gpu/gpu_vulkan.c b/gpu/gpu_vulkan.c new file mode 100644 index 000000000..39058cd75 --- /dev/null +++ b/gpu/gpu_vulkan.c @@ -0,0 +1,126 @@ +#include "gpu_vulkan.h" + +#include + +int check_perfmon() { +#ifdef __linux__ + cap_t caps; + const cap_value_t cap_list[2] = {CAP_PERFMON}; + + if (!CAP_IS_SUPPORTED(CAP_SETFCAP)) + return -1; + + caps = cap_get_proc(); + if (caps == NULL) + return -1; + + if (cap_set_flag(caps, CAP_EFFECTIVE, 2, cap_list, CAP_SET) == -1) + return -1; + + if (cap_set_proc(caps) == -1) + return -1; + + if (cap_free(caps) == -1) + return -1; + + return 0; +#else + return 0; +#endif +} + +void vk_init(vk_init_resp_t *resp) { + if (check_perfmon() != 0) { + resp->err = "Performance monitoring is not allowed. Please enable CAP_PERFMON or run as root to use Vulkan."; + return; + } + + VkInstance instance; + VkApplicationInfo appInfo = {}; + appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO; + appInfo.pNext = NULL; + appInfo.pApplicationName = "Ollama"; + appInfo.applicationVersion = VK_MAKE_VERSION(1, 0, 0); + appInfo.pEngineName = "No Engine"; + appInfo.engineVersion = VK_MAKE_VERSION(1, 0, 0); + appInfo.apiVersion = VK_API_VERSION_1_2; + VkInstanceCreateInfo createInfo = {}; + createInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO; + createInfo.pNext = NULL; + createInfo.flags = 0; + createInfo.enabledExtensionCount = 1; + const char* extensions[] = { VK_KHR_GET_PHYSICAL_DEVICE_PROPERTIES_2_EXTENSION_NAME }; + createInfo.ppEnabledExtensionNames = extensions; + createInfo.pApplicationInfo = &appInfo; + VkResult result = vkCreateInstance(&createInfo, NULL, &instance); + if (result != VK_SUCCESS) { + resp.err = sprintf("Failed to create instance: %d", result); + return; + } + + uint32_t deviceCount; + result = vkEnumeratePhysicalDevices(instance, &deviceCount, NULL); + if (result != VK_SUCCESS) { + resp.err = sprintf("Failed to enumerate physical devices: %d", result); + return; + } + + resp.err = NULL; + resp.oh = instance; + resp.num_devices = deviceCount; +} + +void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp) { + uint32_t deviceCount = rh->num_devices; + VkInstance instance = rh->oh; + + VkPhysicalDevice* devices = malloc(deviceCount * sizeof(VkPhysicalDevice)); + result = vkEnumeratePhysicalDevices(instance, &deviceCount, devices); + if (result != VK_SUCCESS) { + resp.err = sprintf("Failed to enumerate physical devices: %d", result); + return; + } + + VkPhysicalDeviceProperties properties; + vkGetPhysicalDeviceProperties(devices[i], &properties); + LOG(h.verbose, "Vulkan device %d: %s\n", i, properties.deviceName); + int supports_budget = support_memory_budget(devices[i]); + if (!supports_budget) { + resp.err = sprintf("Device %d does not support memory budget\n", i); + return; + } + if (properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_CPU) { + resp.err = sprintf("Device %d is a CPU, skipped\n", i); + return; + } + + VkPhysicalDeviceMemoryBudgetPropertiesEXT physical_device_memory_budget_properties; + physical_device_memory_budget_properties.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_BUDGET_PROPERTIES_EXT; + physical_device_memory_budget_properties.pNext = NULL; + + VkPhysicalDeviceMemoryProperties2 device_memory_properties; + device_memory_properties.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_PROPERTIES_2; + device_memory_properties.pNext = &physical_device_memory_budget_properties; + + vkGetPhysicalDeviceMemoryProperties2(devices[i], &device_memory_properties); + + VkDeviceSize device_memory_total_usage = 0; + VkDeviceSize device_memory_heap_budget = 0; + + for (uint32_t j = 0; j < device_memory_properties.memoryProperties.memoryHeapCount; j++) { + VkMemoryHeap heap = device_memory_properties.memoryProperties.memoryHeaps[j]; + if (heap.flags & VK_MEMORY_HEAP_DEVICE_LOCAL_BIT) { + device_memory_total_usage += physical_device_memory_budget_properties.heapUsage[j]; + device_memory_heap_budget += physical_device_memory_budget_properties.heapBudget[j]; + } + } + + resp->err = NULL; + snprintf(&resp->gpu_id[0], GPU_ID_LEN, "%d", i); + snprintf(&resp->gpu_name[0], GPU_NAME_LEN, "%s", properties.deviceName); + resp->total = (uint64_t) device_memory_total_usage; + resp->free = (uint64_t) device_memory_total_usage; + resp->major = VK_API_VERSION_MAJOR(properties.apiVersion); + resp->minor = VK_API_VERSION_MINOR(properties.apiVersion); + resp->patch = VK_API_VERSION_PATCH(properties.apiVersion); +} diff --git a/gpu/gpu_vulkan.h b/gpu/gpu_vulkan.h new file mode 100644 index 000000000..61ebb1a57 --- /dev/null +++ b/gpu/gpu_vulkan.h @@ -0,0 +1,17 @@ +#include "gpu_info.h" + +#ifdef __linux__ +#include +#endif + +typedef VkInstance vk_handle_t; + +typedef struct vk_init_resp +{ + char *err; // If err is non-null handle is invalid + int num_devices; + vk_handle_t oh; +} vk_init_resp_t; + +void vk_init(vk_init_resp_t *resp); +void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp); From 9c6b0495678f66f5b6b50fdb05c7efd99f5a208f Mon Sep 17 00:00:00 2001 From: pufferffish Date: Sat, 15 Jun 2024 05:27:14 +0100 Subject: [PATCH 02/29] add support in gpu.go --- gpu/gpu.go | 93 +++++++++++++++++++++++++++++++++++++++++++- gpu/gpu_vulkan.c | 24 ++++++++++-- gpu/gpu_vulkan.h | 3 +- gpu/vulkan_common.go | 19 +++++++++ 4 files changed, 133 insertions(+), 6 deletions(-) create mode 100644 gpu/vulkan_common.go diff --git a/gpu/gpu.go b/gpu/gpu.go index a55903c51..359c6b5a7 100644 --- a/gpu/gpu.go +++ b/gpu/gpu.go @@ -29,6 +29,7 @@ type handles struct { cudart *C.cudart_handle_t nvcuda *C.nvcuda_handle_t oneapi *C.oneapi_handle_t + vulkan *C.vk_handle_t } const ( @@ -90,6 +91,16 @@ var OneapiLinuxGlobs = []string{ "/usr/lib*/libze_intel_gpu.so*", } +var VulkanLinuxGlobs = []string{ + "/usr/lib/x86_64-linux-gnu/libvulkan.so*", + "/usr/lib*/libvulkan.so*", +} + +var CapLinuxGlobs = []string{ + "/usr/lib/x86_64-linux-gnu/libcap.so*", + "/usr/lib*/libcap.so*", +} + // Jetson devices have JETSON_JETPACK="x.y.z" factory set to the Jetpack version installed. // Included to drive logic for reducing Ollama-allocated overhead on L4T/Jetson devices. var CudaTegra string = os.Getenv("JETSON_JETPACK") @@ -104,6 +115,10 @@ func initGPUHandles() *handles { var cudartMgmtPatterns []string var nvcudaMgmtName string var nvcudaMgmtPatterns []string + var vulkanMgmtName string + var vulkanMgmtPatterns []string + var libcapMgmtName string + var libcapMgmtPatterns []string tmpDir, _ := PayloadsDir() switch runtime.GOOS { @@ -125,6 +140,12 @@ func initGPUHandles() *handles { // Aligned with driver, we can't carry as payloads nvcudaMgmtName = "libcuda.so*" nvcudaMgmtPatterns = NvcudaLinuxGlobs + + // Vulkan also needs libcap + vulkanMgmtName = "libvulkan.so*" + vulkanMgmtPatterns = VulkanLinuxGlobs + libcapMgmtName = "libcap.so*" + libcapMgmtPatterns = CapLinuxGlobs default: return gpuHandles } @@ -152,6 +173,25 @@ func initGPUHandles() *handles { } } + vulkanLibPaths := FindGPULibs(vulkanMgmtName, vulkanMgmtPatterns) + + var libcapLibPaths []string + if runtime.GOOS == "linux" { + libcapLibPaths = FindGPULibs(libcapMgmtName, libcapMgmtPatterns) + } else { + libcapLibPaths = []string{""} + } + + if len(vulkanLibPaths) > 0 && len(libcapLibPaths) > 0 { + deviceCount, vulkan, vkLibPath, capLibPath := LoadVulkanMgmt(vulkanLibPaths, libcapLibPaths) + if vulkan != nil { + slog.Debug("detected GPUs", "library", vkLibPath, capLibPath, "count", deviceCount) + gpuHandles.vulkan = vulkan + gpuHandles.deviceCount = deviceCount + return gpuHandles + } + } + return gpuHandles } @@ -186,7 +226,7 @@ func GetGPUInfo() GpuInfoList { var memInfo C.mem_info_t resp := []GpuInfo{} - // NVIDIA first + // NVIDIA and Vulkan first for i := range gpuHandles.deviceCount { // TODO once we support CPU compilation variants of GPU libraries refine this... if cpuVariant == "" && runtime.GOARCH == "amd64" { @@ -227,6 +267,32 @@ func GetGPUInfo() GpuInfoList { // TODO potentially sort on our own algorithm instead of what the underlying GPU library does... resp = append(resp, gpuInfo) } + + if gpuHandles.vulkan != nil { + gpuInfo := GpuInfo{ + Library: "vulkan", + } + + C.vk_check_vram(*gpuHandles.vulkan, C.int(i), &memInfo) + if memInfo.err != nil { + slog.Info("error looking up vulkan GPU memory", "error", C.GoString(memInfo.err)) + C.free(unsafe.Pointer(memInfo.err)) + continue + } + + gpuInfo.TotalMemory = uint64(memInfo.total) + gpuInfo.FreeMemory = uint64(memInfo.free) + gpuInfo.ID = C.GoString(&memInfo.gpu_id[0]) + gpuInfo.Compute = fmt.Sprintf("%d.%d", memInfo.major, memInfo.minor) + gpuInfo.MinimumMemory = 0 + gpuInfo.DependencyPath = depPath + gpuInfo.Name = C.GoString(&memInfo.gpu_name[0]) + gpuInfo.DriverMajor = int(memInfo.major) + gpuInfo.DriverMinor = int(memInfo.minor) + + // TODO potentially sort on our own algorithm instead of what the underlying GPU library does... + resp = append(resp, gpuInfo) + } } // Then AMD @@ -379,6 +445,29 @@ func LoadOneapiMgmt(oneapiLibPaths []string) (int, *C.oneapi_handle_t, string) { return 0, nil, "" } +func LoadVulkanMgmt(vulkanLibPaths []string, capLibPaths []string) (int, *C.vk_handle_t, string, string) { + var resp C.vk_init_resp_t + for _, vkLibPath := range vulkanLibPaths { + for _, capLibPath := range capLibPaths { + vkLib := C.CString(vkLibPath) + capLib := C.CString(capLibPath) + defer C.free(unsafe.Pointer(vkLib)) + defer C.free(unsafe.Pointer(capLib)) + + C.vk_init(vkLib, capLib, &resp) + if resp.err != nil { + slog.Debug("Unable to load vulkan", "library", vkLibPath, "error", C.GoString(resp.err)) + slog.Debug("Unable to load libcap", "library", capLibPath, "error", C.GoString(resp.err)) + C.free(unsafe.Pointer(resp.err)) + } else { + return int(resp.num_devices), &resp.vk, vkLibPath, capLibPath + } + } + } + + return 0, nil, "", "" +} + func getVerboseState() C.uint16_t { if envconfig.Debug { return C.uint16_t(1) @@ -401,6 +490,8 @@ func (l GpuInfoList) GetVisibleDevicesEnv() (string, string) { return rocmGetVisibleDevicesEnv(l) case "oneapi": return oneapiGetVisibleDevicesEnv(l) + case "vulkan": + return vkGetVisibleDevicesEnv(l) default: slog.Debug("no filter required for library " + l[0].Library) return "", "" diff --git a/gpu/gpu_vulkan.c b/gpu/gpu_vulkan.c index 39058cd75..bb45bdf21 100644 --- a/gpu/gpu_vulkan.c +++ b/gpu/gpu_vulkan.c @@ -22,18 +22,28 @@ int check_perfmon() { if (cap_free(caps) == -1) return -1; +#endif return 0; -#else - return 0; -#endif } -void vk_init(vk_init_resp_t *resp) { +void vk_init(char* vk_lib_path, char* cap_lib_path, vk_init_resp_t *resp) { + if (!LOAD_LIBRARY(vk_lib_path, RTLD_LAZY)) { + resp->err = "Failed to load Vulkan library"; + return; + } + +#ifdef __linux__ + if (!LOAD_LIBRARY(cap_lib_path, RTLD_LAZY)) { + resp->err = "Failed to load libcap library"; + return; + } + if (check_perfmon() != 0) { resp->err = "Performance monitoring is not allowed. Please enable CAP_PERFMON or run as root to use Vulkan."; return; } +#endif VkInstance instance; VkApplicationInfo appInfo = {}; @@ -123,4 +133,10 @@ void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp) { resp->major = VK_API_VERSION_MAJOR(properties.apiVersion); resp->minor = VK_API_VERSION_MINOR(properties.apiVersion); resp->patch = VK_API_VERSION_PATCH(properties.apiVersion); + } + +void vk_free(vk_handle_t rh) { + vkDestroyInstance(rh->oh, NULL); + free(rh); +} \ No newline at end of file diff --git a/gpu/gpu_vulkan.h b/gpu/gpu_vulkan.h index 61ebb1a57..e77ce554e 100644 --- a/gpu/gpu_vulkan.h +++ b/gpu/gpu_vulkan.h @@ -13,5 +13,6 @@ typedef struct vk_init_resp vk_handle_t oh; } vk_init_resp_t; -void vk_init(vk_init_resp_t *resp); +void vk_init(char* vk_lib_path, char* cap_lib_path, vk_init_resp_t *resp); void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp); +void vk_free(vk_handle_t rh); diff --git a/gpu/vulkan_common.go b/gpu/vulkan_common.go new file mode 100644 index 000000000..8d3d15d06 --- /dev/null +++ b/gpu/vulkan_common.go @@ -0,0 +1,19 @@ +package gpu + +import ( + "log/slog" + "strings" +) + +func vkGetVisibleDevicesEnv(gpuInfo []GpuInfo) (string, string) { + ids := []string{} + for _, info := range gpuInfo { + if info.Library != "vulkan" { + // TODO shouldn't happen if things are wired correctly... + slog.Debug("vkGetVisibleDevicesEnv skipping over non-vulkan device", "library", info.Library) + continue + } + ids = append(ids, info.ID) + } + return "GGML_VK_VISIBLE_DEVICES", strings.Join(ids, ",") +} From 93c4d69daa02be2c4407c73d30c8fe72961de61b Mon Sep 17 00:00:00 2001 From: pufferffish Date: Sat, 15 Jun 2024 05:42:59 +0100 Subject: [PATCH 03/29] add support in gen_linux.sh --- llm/generate/gen_linux.sh | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/llm/generate/gen_linux.sh b/llm/generate/gen_linux.sh index a9df6ff86..2190fb93e 100755 --- a/llm/generate/gen_linux.sh +++ b/llm/generate/gen_linux.sh @@ -206,6 +206,34 @@ if [ -z "${OLLAMA_SKIP_CUDA_GENERATE}" -a -d "${CUDA_LIB_DIR}" ]; then fi +if [ -z "${VULKAN_ROOT}" ]; then + # Try the default location in case it exists + VULKAN_ROOT=/usr/lib/ +fi + +if [ -z "${CAP_ROOT}" ]; then + # Try the default location in case it exists + CAP_ROOT=/usr/lib/ +fi + +if [ -z "${OLLAMA_SKIP_VULKAN_GENERATE}" -a -d "${VULKAN_ROOT}" ] && [ -z "${OLLAMA_SKIP_VULKAN_GENERATE}" -a -d "${CAP_ROOT}" ]; then + echo "Vulkan and capabilities libraries detected - building dynamic Vulkan library" + init_vars + + CMAKE_DEFS="${COMMON_CMAKE_DEFS} ${CMAKE_DEFS} -DLLAMA_VULKAN=1" + BUILD_DIR="../build/linux/${ARCH}/vulkan" + EXTRA_LIBS="-L${VULKAN_ROOT} -L${CAP_ROOT} -lvulkan -lcap" + build + + # copy oneAPI dependencies + for dep in $(ldd "${BUILD_DIR}/bin/ollama_llama_server" | grep "=>" | cut -f2 -d= | cut -f2 -d' ' | grep -e vulkan -e cap); do + cp "${dep}" "${BUILD_DIR}/bin/" + done + cp "${VULKAN_ROOT}/libvulkan.so" "${BUILD_DIR}/bin/" + cp "${CAP_ROOT}/libcap.so" "${BUILD_DIR}/bin/" + compress +fi + if [ -z "${ONEAPI_ROOT}" ]; then # Try the default location in case it exists ONEAPI_ROOT=/opt/intel/oneapi From 24c8840037a9edd48fafd31f113916cb4105c922 Mon Sep 17 00:00:00 2001 From: pufferffish Date: Sat, 15 Jun 2024 07:49:28 +0100 Subject: [PATCH 04/29] it builds --- gpu/gpu.go | 6 +- gpu/gpu_info.h | 1 + gpu/gpu_info_vulkan.c | 222 ++++++++++++++++++++++++++++++++++++++++++ gpu/gpu_info_vulkan.h | 66 +++++++++++++ gpu/gpu_vulkan.c | 142 --------------------------- gpu/gpu_vulkan.h | 18 ---- 6 files changed, 294 insertions(+), 161 deletions(-) create mode 100644 gpu/gpu_info_vulkan.c create mode 100644 gpu/gpu_info_vulkan.h delete mode 100644 gpu/gpu_vulkan.c delete mode 100644 gpu/gpu_vulkan.h diff --git a/gpu/gpu.go b/gpu/gpu.go index 359c6b5a7..0b19e0aba 100644 --- a/gpu/gpu.go +++ b/gpu/gpu.go @@ -209,6 +209,9 @@ func GetGPUInfo() GpuInfoList { if gpuHandles.nvcuda != nil { C.nvcuda_release(*gpuHandles.nvcuda) } + if gpuHandles.vulkan != nil { + C.vk_release(*gpuHandles.vulkan) + } }() // All our GPU builds on x86 have AVX enabled, so fallback to CPU if we don't detect at least AVX @@ -447,6 +450,7 @@ func LoadOneapiMgmt(oneapiLibPaths []string) (int, *C.oneapi_handle_t, string) { func LoadVulkanMgmt(vulkanLibPaths []string, capLibPaths []string) (int, *C.vk_handle_t, string, string) { var resp C.vk_init_resp_t + resp.ch.verbose = getVerboseState() for _, vkLibPath := range vulkanLibPaths { for _, capLibPath := range capLibPaths { vkLib := C.CString(vkLibPath) @@ -460,7 +464,7 @@ func LoadVulkanMgmt(vulkanLibPaths []string, capLibPaths []string) (int, *C.vk_h slog.Debug("Unable to load libcap", "library", capLibPath, "error", C.GoString(resp.err)) C.free(unsafe.Pointer(resp.err)) } else { - return int(resp.num_devices), &resp.vk, vkLibPath, capLibPath + return int(resp.num_devices), &resp.ch, vkLibPath, capLibPath } } } diff --git a/gpu/gpu_info.h b/gpu/gpu_info.h index 482b81a6f..afc80dca4 100644 --- a/gpu/gpu_info.h +++ b/gpu/gpu_info.h @@ -63,6 +63,7 @@ void cpu_check_ram(mem_info_t *resp); #include "gpu_info_cudart.h" #include "gpu_info_nvcuda.h" #include "gpu_info_oneapi.h" +#include "gpu_info_vulkan.h" #endif // __GPU_INFO_H__ #endif // __APPLE__ \ No newline at end of file diff --git a/gpu/gpu_info_vulkan.c b/gpu/gpu_info_vulkan.c new file mode 100644 index 000000000..8b0370d2c --- /dev/null +++ b/gpu/gpu_info_vulkan.c @@ -0,0 +1,222 @@ +#include "gpu_info_vulkan.h" + +#include + +int check_perfmon(vk_handle_t* rh) { +#ifdef __linux__ + cap_t caps; + const cap_value_t cap_list[2] = {CAP_PERFMON}; + + if ((*rh->cap_get_bound)(CAP_SETFCAP) < 0) + return -1; + + caps = (*rh->cap_get_proc)(); + if (caps == NULL) + return -1; + + if ((*rh->cap_set_flag)(caps, CAP_EFFECTIVE, 2, cap_list, CAP_SET) == -1) + return -1; + + if ((*rh->cap_set_proc)(caps) == -1) + return -1; + + if ((*rh->cap_free)(caps) == -1) + return -1; +#endif + + return 0; +} + +int support_memory_budget(vk_handle_t* rh, VkPhysicalDevice device) { + VkPhysicalDeviceProperties properties; + (*rh->vkGetPhysicalDeviceProperties)(device, &properties); + uint32_t extensionCount; + (*rh->vkEnumerateDeviceExtensionProperties)(device, NULL, &extensionCount, NULL); + VkExtensionProperties* extensions = malloc(extensionCount * sizeof(VkExtensionProperties)); + (*rh->vkEnumerateDeviceExtensionProperties)(device, NULL, &extensionCount, extensions); + for (int j = 0; j < extensionCount; j++) { + if (strcmp(extensions[j].extensionName, VK_EXT_MEMORY_BUDGET_EXTENSION_NAME) == 0) { + return 1; + } + } + return 0; +} + +void vk_init(char* vk_lib_path, char* cap_lib_path, vk_init_resp_t *resp) { + const int buflen = 256; + char buf[buflen + 1]; + int i; + + struct lookup { + int is_cap; + char *s; + void **p; + } l[] = { +#ifdef __linux__ + {1, "cap_get_proc", (void *)&resp->ch.cap_get_proc}, + {1, "cap_get_bound", (void *)&resp->ch.cap_get_bound}, + {1, "cap_set_flag", (void *)&resp->ch.cap_set_flag}, + {1, "cap_set_proc", (void *)&resp->ch.cap_set_proc}, + {1, "cap_free", (void *)&resp->ch.cap_free}, +#endif + {0, "vkGetPhysicalDeviceProperties", (void *)&resp->ch.vkGetPhysicalDeviceProperties}, + {0, "vkEnumerateDeviceExtensionProperties", (void *)&resp->ch.vkEnumerateDeviceExtensionProperties}, + {0, "vkCreateInstance", (void *)&resp->ch.vkCreateInstance}, + {0, "vkEnumeratePhysicalDevices", (void *)&resp->ch.vkEnumeratePhysicalDevices}, + {0, "vkGetPhysicalDeviceMemoryProperties2", (void *)&resp->ch.vkGetPhysicalDeviceMemoryProperties2}, + {0, "vkDestroyInstance", (void *)&resp->ch.vkDestroyInstance}, + {0, NULL, NULL}, + }; + + resp->ch.vk_handle = LOAD_LIBRARY(vk_lib_path, RTLD_LAZY); + if (!resp->ch.vk_handle) { + char *msg = LOAD_ERR(); + LOG(resp->ch.verbose, "library %s load err: %s\n", vk_lib_path, msg); + snprintf(buf, buflen, + "Unable to load %s library to query for Vulkan GPUs: %s", + vk_lib_path, msg); + free(msg); + resp->err = strdup(buf); + return; + } + +#ifdef __linux__ + resp->ch.cap_handle = LOAD_LIBRARY(cap_lib_path, RTLD_LAZY); + if (!resp->ch.cap_handle) { + char *msg = LOAD_ERR(); + LOG(resp->ch.verbose, "library %s load err: %s\n", cap_lib_path, msg); + snprintf(buf, buflen, + "Unable to load %s library to query for Vulkan GPUs: %s", + cap_lib_path, msg); + free(msg); + resp->err = strdup(buf); + return; + } +#endif + + for (i = 0; l[i].s != NULL; i++) { + if (l[i].is_cap) +#ifdef __linux__ + *l[i].p = LOAD_SYMBOL(resp->ch.cap_handle, l[i].s); +#else + continue; +#endif + else + *l[i].p = LOAD_SYMBOL(resp->ch.vk_handle, l[i].s); + if (!*l[i].p) { + char *msg = LOAD_ERR(); + LOG(resp->ch.verbose, "dlerr: %s\n", msg); + if (l[i].is_cap) { + UNLOAD_LIBRARY(resp->ch.cap_handle); + resp->ch.cap_handle = NULL; + } else { + UNLOAD_LIBRARY(resp->ch.vk_handle); + resp->ch.vk_handle = NULL; + } + snprintf(buf, buflen, "symbol lookup for %s failed: %s", l[i].s, + msg); + free(msg); + resp->err = strdup(buf); + return; + } + } + + if (check_perfmon(&resp->ch) != 0) { + resp->err = "Performance monitoring is not allowed. Please enable CAP_PERFMON or run as root to use Vulkan."; + return; + } + + VkInstance instance; + VkApplicationInfo appInfo = {}; + appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO; + appInfo.pNext = NULL; + appInfo.pApplicationName = "Ollama"; + appInfo.applicationVersion = VK_MAKE_VERSION(1, 0, 0); + appInfo.pEngineName = "No Engine"; + appInfo.engineVersion = VK_MAKE_VERSION(1, 0, 0); + appInfo.apiVersion = VK_API_VERSION_1_2; + VkInstanceCreateInfo createInfo = {}; + createInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO; + createInfo.pNext = NULL; + createInfo.flags = 0; + createInfo.enabledExtensionCount = 1; + const char* extensions[] = { VK_KHR_GET_PHYSICAL_DEVICE_PROPERTIES_2_EXTENSION_NAME }; + createInfo.ppEnabledExtensionNames = extensions; + createInfo.pApplicationInfo = &appInfo; + VkResult result = (*resp->ch.vkCreateInstance)(&createInfo, NULL, &instance); + if (result != VK_SUCCESS) { + resp->err = strdup("failed to create instance"); + return; + } + + uint32_t deviceCount; + result = (*resp->ch.vkEnumeratePhysicalDevices)(instance, &deviceCount, NULL); + if (result != VK_SUCCESS) { + resp->err = strdup("failed to enumerate physical devices"); + return; + } + + resp->err = NULL; + resp->ch.vk = instance; + resp->ch.num_devices = deviceCount; + resp->num_devices = deviceCount; +} + +void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp) { + VkInstance instance = rh.vk; + uint32_t deviceCount = rh.num_devices; + + VkPhysicalDevice* devices = malloc(deviceCount * sizeof(VkPhysicalDevice)); + VkResult result = (*rh.vkEnumeratePhysicalDevices)(instance, &deviceCount, devices); + if (result != VK_SUCCESS) { + resp->err = strdup("failed to enumerate physical devices"); + return; + } + + VkPhysicalDeviceProperties properties; + (*rh.vkGetPhysicalDeviceProperties)(devices[i], &properties); + int supports_budget = support_memory_budget(&rh, devices[i]); + if (!supports_budget) { + resp->err = strdup("device does not support memory budget"); + return; + } + if (properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_CPU) { + resp->err = strdup("device is a CPU"); + return; + } + + VkPhysicalDeviceMemoryBudgetPropertiesEXT physical_device_memory_budget_properties; + physical_device_memory_budget_properties.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_BUDGET_PROPERTIES_EXT; + physical_device_memory_budget_properties.pNext = NULL; + + VkPhysicalDeviceMemoryProperties2 device_memory_properties; + device_memory_properties.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_PROPERTIES_2; + device_memory_properties.pNext = &physical_device_memory_budget_properties; + + (*rh.vkGetPhysicalDeviceMemoryProperties2)(devices[i], &device_memory_properties); + + VkDeviceSize device_memory_total_usage = 0; + VkDeviceSize device_memory_heap_budget = 0; + + for (uint32_t j = 0; j < device_memory_properties.memoryProperties.memoryHeapCount; j++) { + VkMemoryHeap heap = device_memory_properties.memoryProperties.memoryHeaps[j]; + if (heap.flags & VK_MEMORY_HEAP_DEVICE_LOCAL_BIT) { + device_memory_total_usage += physical_device_memory_budget_properties.heapUsage[j]; + device_memory_heap_budget += physical_device_memory_budget_properties.heapBudget[j]; + } + } + + resp->err = NULL; + snprintf(&resp->gpu_id[0], GPU_ID_LEN, "%d", i); + snprintf(&resp->gpu_name[0], GPU_NAME_LEN, "%s", properties.deviceName); + resp->total = (uint64_t) device_memory_total_usage; + resp->free = (uint64_t) device_memory_total_usage; + resp->major = VK_API_VERSION_MAJOR(properties.apiVersion); + resp->minor = VK_API_VERSION_MINOR(properties.apiVersion); + resp->patch = VK_API_VERSION_PATCH(properties.apiVersion); + +} + +void vk_release(vk_handle_t rh) { + (*rh.vkDestroyInstance)(rh.vk, NULL); +} \ No newline at end of file diff --git a/gpu/gpu_info_vulkan.h b/gpu/gpu_info_vulkan.h new file mode 100644 index 000000000..6025f3e09 --- /dev/null +++ b/gpu/gpu_info_vulkan.h @@ -0,0 +1,66 @@ +#ifndef __APPLE__ +#ifndef __GPU_INFO_VULKAN_H__ +#define __GPU_INFO_VULKAN_H__ + +#include "gpu_info.h" + +#ifdef __linux__ +#include +#endif + +#include + +typedef struct { + void* vk_handle; + void* cap_handle; + uint16_t verbose; + + VkInstance vk; + int num_devices; + +#ifdef __linux__ + cap_t (*cap_get_proc)(void); + + int (*cap_get_bound)(cap_value_t); + int (*cap_set_flag)(cap_t, cap_flag_t, int, const cap_value_t *, cap_flag_value_t); + int (*cap_set_proc)(cap_t); + int (*cap_free)(cap_t); +#endif + + void (*vkGetPhysicalDeviceProperties)( + VkPhysicalDevice physicalDevice, + VkPhysicalDeviceProperties* pProperties); + VkResult (*vkEnumerateDeviceExtensionProperties)( + VkPhysicalDevice physicalDevice, + const char* pLayerName, + uint32_t* pPropertyCount, + VkExtensionProperties* pProperties); + VkResult (*vkCreateInstance)( + const VkInstanceCreateInfo* pCreateInfo, + const VkAllocationCallbacks* pAllocator, + VkInstance* pInstance); + VkResult (*vkEnumeratePhysicalDevices)( + VkInstance instance, + uint32_t* pPhysicalDeviceCount, + VkPhysicalDevice* pPhysicalDevices); + void (*vkGetPhysicalDeviceMemoryProperties2)( + VkPhysicalDevice physicalDevice, + VkPhysicalDeviceMemoryProperties2* pMemoryProperties); + void (*vkDestroyInstance)( + VkInstance instance, + const VkAllocationCallbacks* pAllocator); +} vk_handle_t; + +typedef struct vk_init_resp +{ + char *err; // If err is non-null handle is invalid + int num_devices; + vk_handle_t ch; +} vk_init_resp_t; + +void vk_init(char* vk_lib_path, char* cap_lib_path, vk_init_resp_t *resp); +void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp); +void vk_release(vk_handle_t rh); + +#endif +#endif \ No newline at end of file diff --git a/gpu/gpu_vulkan.c b/gpu/gpu_vulkan.c deleted file mode 100644 index bb45bdf21..000000000 --- a/gpu/gpu_vulkan.c +++ /dev/null @@ -1,142 +0,0 @@ -#include "gpu_vulkan.h" - -#include - -int check_perfmon() { -#ifdef __linux__ - cap_t caps; - const cap_value_t cap_list[2] = {CAP_PERFMON}; - - if (!CAP_IS_SUPPORTED(CAP_SETFCAP)) - return -1; - - caps = cap_get_proc(); - if (caps == NULL) - return -1; - - if (cap_set_flag(caps, CAP_EFFECTIVE, 2, cap_list, CAP_SET) == -1) - return -1; - - if (cap_set_proc(caps) == -1) - return -1; - - if (cap_free(caps) == -1) - return -1; -#endif - - return 0; -} - -void vk_init(char* vk_lib_path, char* cap_lib_path, vk_init_resp_t *resp) { - if (!LOAD_LIBRARY(vk_lib_path, RTLD_LAZY)) { - resp->err = "Failed to load Vulkan library"; - return; - } - -#ifdef __linux__ - if (!LOAD_LIBRARY(cap_lib_path, RTLD_LAZY)) { - resp->err = "Failed to load libcap library"; - return; - } - - if (check_perfmon() != 0) { - resp->err = "Performance monitoring is not allowed. Please enable CAP_PERFMON or run as root to use Vulkan."; - return; - } -#endif - - VkInstance instance; - VkApplicationInfo appInfo = {}; - appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO; - appInfo.pNext = NULL; - appInfo.pApplicationName = "Ollama"; - appInfo.applicationVersion = VK_MAKE_VERSION(1, 0, 0); - appInfo.pEngineName = "No Engine"; - appInfo.engineVersion = VK_MAKE_VERSION(1, 0, 0); - appInfo.apiVersion = VK_API_VERSION_1_2; - VkInstanceCreateInfo createInfo = {}; - createInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO; - createInfo.pNext = NULL; - createInfo.flags = 0; - createInfo.enabledExtensionCount = 1; - const char* extensions[] = { VK_KHR_GET_PHYSICAL_DEVICE_PROPERTIES_2_EXTENSION_NAME }; - createInfo.ppEnabledExtensionNames = extensions; - createInfo.pApplicationInfo = &appInfo; - VkResult result = vkCreateInstance(&createInfo, NULL, &instance); - if (result != VK_SUCCESS) { - resp.err = sprintf("Failed to create instance: %d", result); - return; - } - - uint32_t deviceCount; - result = vkEnumeratePhysicalDevices(instance, &deviceCount, NULL); - if (result != VK_SUCCESS) { - resp.err = sprintf("Failed to enumerate physical devices: %d", result); - return; - } - - resp.err = NULL; - resp.oh = instance; - resp.num_devices = deviceCount; -} - -void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp) { - uint32_t deviceCount = rh->num_devices; - VkInstance instance = rh->oh; - - VkPhysicalDevice* devices = malloc(deviceCount * sizeof(VkPhysicalDevice)); - result = vkEnumeratePhysicalDevices(instance, &deviceCount, devices); - if (result != VK_SUCCESS) { - resp.err = sprintf("Failed to enumerate physical devices: %d", result); - return; - } - - VkPhysicalDeviceProperties properties; - vkGetPhysicalDeviceProperties(devices[i], &properties); - LOG(h.verbose, "Vulkan device %d: %s\n", i, properties.deviceName); - int supports_budget = support_memory_budget(devices[i]); - if (!supports_budget) { - resp.err = sprintf("Device %d does not support memory budget\n", i); - return; - } - if (properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_CPU) { - resp.err = sprintf("Device %d is a CPU, skipped\n", i); - return; - } - - VkPhysicalDeviceMemoryBudgetPropertiesEXT physical_device_memory_budget_properties; - physical_device_memory_budget_properties.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_BUDGET_PROPERTIES_EXT; - physical_device_memory_budget_properties.pNext = NULL; - - VkPhysicalDeviceMemoryProperties2 device_memory_properties; - device_memory_properties.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_PROPERTIES_2; - device_memory_properties.pNext = &physical_device_memory_budget_properties; - - vkGetPhysicalDeviceMemoryProperties2(devices[i], &device_memory_properties); - - VkDeviceSize device_memory_total_usage = 0; - VkDeviceSize device_memory_heap_budget = 0; - - for (uint32_t j = 0; j < device_memory_properties.memoryProperties.memoryHeapCount; j++) { - VkMemoryHeap heap = device_memory_properties.memoryProperties.memoryHeaps[j]; - if (heap.flags & VK_MEMORY_HEAP_DEVICE_LOCAL_BIT) { - device_memory_total_usage += physical_device_memory_budget_properties.heapUsage[j]; - device_memory_heap_budget += physical_device_memory_budget_properties.heapBudget[j]; - } - } - - resp->err = NULL; - snprintf(&resp->gpu_id[0], GPU_ID_LEN, "%d", i); - snprintf(&resp->gpu_name[0], GPU_NAME_LEN, "%s", properties.deviceName); - resp->total = (uint64_t) device_memory_total_usage; - resp->free = (uint64_t) device_memory_total_usage; - resp->major = VK_API_VERSION_MAJOR(properties.apiVersion); - resp->minor = VK_API_VERSION_MINOR(properties.apiVersion); - resp->patch = VK_API_VERSION_PATCH(properties.apiVersion); - -} - -void vk_free(vk_handle_t rh) { - vkDestroyInstance(rh->oh, NULL); - free(rh); -} \ No newline at end of file diff --git a/gpu/gpu_vulkan.h b/gpu/gpu_vulkan.h deleted file mode 100644 index e77ce554e..000000000 --- a/gpu/gpu_vulkan.h +++ /dev/null @@ -1,18 +0,0 @@ -#include "gpu_info.h" - -#ifdef __linux__ -#include -#endif - -typedef VkInstance vk_handle_t; - -typedef struct vk_init_resp -{ - char *err; // If err is non-null handle is invalid - int num_devices; - vk_handle_t oh; -} vk_init_resp_t; - -void vk_init(char* vk_lib_path, char* cap_lib_path, vk_init_resp_t *resp); -void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp); -void vk_free(vk_handle_t rh); From 724fac470f0df86e8d0d24e209bea34f31a4ec84 Mon Sep 17 00:00:00 2001 From: pufferffish Date: Sat, 15 Jun 2024 08:05:48 +0100 Subject: [PATCH 05/29] fix segfault --- gpu/gpu.go | 3 +-- gpu/gpu_info_vulkan.c | 3 ++- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/gpu/gpu.go b/gpu/gpu.go index 0b19e0aba..46359e340 100644 --- a/gpu/gpu.go +++ b/gpu/gpu.go @@ -460,8 +460,7 @@ func LoadVulkanMgmt(vulkanLibPaths []string, capLibPaths []string) (int, *C.vk_h C.vk_init(vkLib, capLib, &resp) if resp.err != nil { - slog.Debug("Unable to load vulkan", "library", vkLibPath, "error", C.GoString(resp.err)) - slog.Debug("Unable to load libcap", "library", capLibPath, "error", C.GoString(resp.err)) + slog.Debug("Unable to load vulkan", "library", vkLibPath, capLibPath, "error", C.GoString(resp.err)) C.free(unsafe.Pointer(resp.err)) } else { return int(resp.num_devices), &resp.ch, vkLibPath, capLibPath diff --git a/gpu/gpu_info_vulkan.c b/gpu/gpu_info_vulkan.c index 8b0370d2c..cb2e8f67e 100644 --- a/gpu/gpu_info_vulkan.c +++ b/gpu/gpu_info_vulkan.c @@ -122,7 +122,8 @@ void vk_init(char* vk_lib_path, char* cap_lib_path, vk_init_resp_t *resp) { } if (check_perfmon(&resp->ch) != 0) { - resp->err = "Performance monitoring is not allowed. Please enable CAP_PERFMON or run as root to use Vulkan."; + resp->err = strdup("performance monitoring is not allowed. Please enable CAP_PERFMON or run as root to use Vulkan."); + LOG(resp->ch.verbose, resp->err); return; } From e4e8a5d25a375c9df03ad122211237798e4ca743 Mon Sep 17 00:00:00 2001 From: pufferffish Date: Sat, 15 Jun 2024 09:44:10 +0100 Subject: [PATCH 06/29] fix compilation --- gpu/gpu_info_vulkan.c | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/gpu/gpu_info_vulkan.c b/gpu/gpu_info_vulkan.c index cb2e8f67e..9822a63f9 100644 --- a/gpu/gpu_info_vulkan.c +++ b/gpu/gpu_info_vulkan.c @@ -123,7 +123,7 @@ void vk_init(char* vk_lib_path, char* cap_lib_path, vk_init_resp_t *resp) { if (check_perfmon(&resp->ch) != 0) { resp->err = strdup("performance monitoring is not allowed. Please enable CAP_PERFMON or run as root to use Vulkan."); - LOG(resp->ch.verbose, resp->err); + LOG(resp->ch.verbose, "vulkan: %s", resp->err); return; } @@ -209,7 +209,8 @@ void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp) { resp->err = NULL; snprintf(&resp->gpu_id[0], GPU_ID_LEN, "%d", i); - snprintf(&resp->gpu_name[0], GPU_NAME_LEN, "%s", properties.deviceName); + resp->gpu_name[GPU_NAME_LEN - 1] = '\0'; + strncpy(&resp->gpu_name[0], properties.deviceName, GPU_NAME_LEN - 1); resp->total = (uint64_t) device_memory_total_usage; resp->free = (uint64_t) device_memory_total_usage; resp->major = VK_API_VERSION_MAJOR(properties.apiVersion); From 257364cb3c47a5e392bfb1772ecf6709dc0a7c83 Mon Sep 17 00:00:00 2001 From: pufferffish Date: Sat, 15 Jun 2024 10:52:34 +0100 Subject: [PATCH 07/29] fix free memory monitor --- gpu/gpu_info_vulkan.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gpu/gpu_info_vulkan.c b/gpu/gpu_info_vulkan.c index 9822a63f9..b4b7f26fd 100644 --- a/gpu/gpu_info_vulkan.c +++ b/gpu/gpu_info_vulkan.c @@ -212,7 +212,7 @@ void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp) { resp->gpu_name[GPU_NAME_LEN - 1] = '\0'; strncpy(&resp->gpu_name[0], properties.deviceName, GPU_NAME_LEN - 1); resp->total = (uint64_t) device_memory_total_usage; - resp->free = (uint64_t) device_memory_total_usage; + resp->free = (uint64_t) device_memory_heap_budget; resp->major = VK_API_VERSION_MAJOR(properties.apiVersion); resp->minor = VK_API_VERSION_MINOR(properties.apiVersion); resp->patch = VK_API_VERSION_PATCH(properties.apiVersion); From 11c55fab8113a02fbd77968c99856c22fb89c880 Mon Sep 17 00:00:00 2001 From: pufferffish Date: Sat, 15 Jun 2024 10:58:12 +0100 Subject: [PATCH 08/29] fix total memory monitor --- gpu/gpu_info_vulkan.c | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gpu/gpu_info_vulkan.c b/gpu/gpu_info_vulkan.c index b4b7f26fd..fbe7a5885 100644 --- a/gpu/gpu_info_vulkan.c +++ b/gpu/gpu_info_vulkan.c @@ -196,13 +196,13 @@ void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp) { (*rh.vkGetPhysicalDeviceMemoryProperties2)(devices[i], &device_memory_properties); - VkDeviceSize device_memory_total_usage = 0; + VkDeviceSize device_memory_total_size = 0; VkDeviceSize device_memory_heap_budget = 0; for (uint32_t j = 0; j < device_memory_properties.memoryProperties.memoryHeapCount; j++) { VkMemoryHeap heap = device_memory_properties.memoryProperties.memoryHeaps[j]; if (heap.flags & VK_MEMORY_HEAP_DEVICE_LOCAL_BIT) { - device_memory_total_usage += physical_device_memory_budget_properties.heapUsage[j]; + device_memory_total_size += heap.size; device_memory_heap_budget += physical_device_memory_budget_properties.heapBudget[j]; } } @@ -211,7 +211,7 @@ void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp) { snprintf(&resp->gpu_id[0], GPU_ID_LEN, "%d", i); resp->gpu_name[GPU_NAME_LEN - 1] = '\0'; strncpy(&resp->gpu_name[0], properties.deviceName, GPU_NAME_LEN - 1); - resp->total = (uint64_t) device_memory_total_usage; + resp->total = (uint64_t) device_memory_total_size; resp->free = (uint64_t) device_memory_heap_budget; resp->major = VK_API_VERSION_MAJOR(properties.apiVersion); resp->minor = VK_API_VERSION_MINOR(properties.apiVersion); From 18f3f960b01e1dd18a43fbcddbc0dc9de1ae2cbd Mon Sep 17 00:00:00 2001 From: pufferffish Date: Sat, 15 Jun 2024 12:05:01 +0100 Subject: [PATCH 09/29] update gpu.go --- gpu/gpu.go | 642 +++++++++++++++++++++++++++++---------------- gpu/gpu_linux.go | 16 ++ gpu/gpu_windows.go | 4 + gpu/types.go | 7 + 4 files changed, 441 insertions(+), 228 deletions(-) diff --git a/gpu/gpu.go b/gpu/gpu.go index 46359e340..11c72e151 100644 --- a/gpu/gpu.go +++ b/gpu/gpu.go @@ -24,20 +24,45 @@ import ( "github.com/ollama/ollama/format" ) -type handles struct { +type cudaHandles struct { deviceCount int cudart *C.cudart_handle_t nvcuda *C.nvcuda_handle_t + nvml *C.nvml_handle_t +} + +type oneapiHandles struct { oneapi *C.oneapi_handle_t - vulkan *C.vk_handle_t + deviceCount int +} + +type vulkanHandles struct { + vulkan *C.vulkan_handle_t + deviceCount int } const ( cudaMinimumMemory = 457 * format.MebiByte rocmMinimumMemory = 457 * format.MebiByte + // TODO OneAPI minimum memory ) -var gpuMutex sync.Mutex +var ( + gpuMutex sync.Mutex + bootstrapped bool + cpuCapability CPUCapability + cpus []CPUInfo + cudaGPUs []CudaGPUInfo + nvcudaLibPath string + cudartLibPath string + oneapiLibPath string + vulkanLibPath string + libcapLibPath string + nvmlLibPath string + rocmGPUs []RocmGPUInfo + oneapiGPUs []OneapiGPUInfo + vulkanGPUs []VulkanGPUInfo +) // With our current CUDA compile flags, older than 5.0 will not work properly var CudaComputeMin = [2]C.int{5, 0} @@ -47,152 +72,133 @@ var RocmComputeMin = 9 // TODO find a better way to detect iGPU instead of minimum memory const IGPUMemLimit = 1 * format.GibiByte // 512G is what they typically report, so anything less than 1G must be iGPU -var CudartLinuxGlobs = []string{ - "/usr/local/cuda/lib64/libcudart.so*", - "/usr/lib/x86_64-linux-gnu/nvidia/current/libcudart.so*", - "/usr/lib/x86_64-linux-gnu/libcudart.so*", - "/usr/lib/wsl/lib/libcudart.so*", - "/usr/lib/wsl/drivers/*/libcudart.so*", - "/opt/cuda/lib64/libcudart.so*", - "/usr/local/cuda*/targets/aarch64-linux/lib/libcudart.so*", - "/usr/lib/aarch64-linux-gnu/nvidia/current/libcudart.so*", - "/usr/lib/aarch64-linux-gnu/libcudart.so*", - "/usr/local/cuda/lib*/libcudart.so*", - "/usr/lib*/libcudart.so*", - "/usr/local/lib*/libcudart.so*", -} - -var CudartWindowsGlobs = []string{ - "c:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v*\\bin\\cudart64_*.dll", -} - -var NvcudaLinuxGlobs = []string{ - "/usr/local/cuda*/targets/*/lib/libcuda.so*", - "/usr/lib/*-linux-gnu/nvidia/current/libcuda.so*", - "/usr/lib/*-linux-gnu/libcuda.so*", - "/usr/lib/wsl/lib/libcuda.so*", - "/usr/lib/wsl/drivers/*/libcuda.so*", - "/opt/cuda/lib*/libcuda.so*", - "/usr/local/cuda/lib*/libcuda.so*", - "/usr/lib*/libcuda.so*", - "/usr/local/lib*/libcuda.so*", -} - -var NvcudaWindowsGlobs = []string{ - "c:\\windows\\system*\\nvcuda.dll", -} - -var OneapiWindowsGlobs = []string{ - "c:\\Windows\\System32\\DriverStore\\FileRepository\\*\\ze_intel_gpu64.dll", -} - -var OneapiLinuxGlobs = []string{ - "/usr/lib/x86_64-linux-gnu/libze_intel_gpu.so*", - "/usr/lib*/libze_intel_gpu.so*", -} - -var VulkanLinuxGlobs = []string{ - "/usr/lib/x86_64-linux-gnu/libvulkan.so*", - "/usr/lib*/libvulkan.so*", -} - -var CapLinuxGlobs = []string{ - "/usr/lib/x86_64-linux-gnu/libcap.so*", - "/usr/lib*/libcap.so*", -} - // Jetson devices have JETSON_JETPACK="x.y.z" factory set to the Jetpack version installed. // Included to drive logic for reducing Ollama-allocated overhead on L4T/Jetson devices. var CudaTegra string = os.Getenv("JETSON_JETPACK") // Note: gpuMutex must already be held -func initGPUHandles() *handles { +func initCudaHandles() *cudaHandles { // TODO - if the ollama build is CPU only, don't do these checks as they're irrelevant and confusing - gpuHandles := &handles{} - var cudartMgmtName string - var cudartMgmtPatterns []string - var nvcudaMgmtName string - var nvcudaMgmtPatterns []string - var vulkanMgmtName string - var vulkanMgmtPatterns []string - var libcapMgmtName string - var libcapMgmtPatterns []string - - tmpDir, _ := PayloadsDir() - switch runtime.GOOS { - case "windows": - cudartMgmtName = "cudart64_*.dll" - localAppData := os.Getenv("LOCALAPPDATA") - cudartMgmtPatterns = []string{filepath.Join(localAppData, "Programs", "Ollama", cudartMgmtName)} - cudartMgmtPatterns = append(cudartMgmtPatterns, CudartWindowsGlobs...) - // Aligned with driver, we can't carry as payloads - nvcudaMgmtName = "nvcuda.dll" - nvcudaMgmtPatterns = NvcudaWindowsGlobs - case "linux": - cudartMgmtName = "libcudart.so*" - if tmpDir != "" { - // TODO - add "payloads" for subprocess - cudartMgmtPatterns = []string{filepath.Join(tmpDir, "cuda*", cudartMgmtName)} - } - cudartMgmtPatterns = append(cudartMgmtPatterns, CudartLinuxGlobs...) - // Aligned with driver, we can't carry as payloads - nvcudaMgmtName = "libcuda.so*" - nvcudaMgmtPatterns = NvcudaLinuxGlobs - - // Vulkan also needs libcap - vulkanMgmtName = "libvulkan.so*" - vulkanMgmtPatterns = VulkanLinuxGlobs - libcapMgmtName = "libcap.so*" - libcapMgmtPatterns = CapLinuxGlobs - default: - return gpuHandles + cHandles := &cudaHandles{} + // Short Circuit if we already know which library to use + if nvmlLibPath != "" { + cHandles.nvml, _ = LoadNVMLMgmt([]string{nvmlLibPath}) + return cHandles + } + if nvcudaLibPath != "" { + cHandles.deviceCount, cHandles.nvcuda, _ = LoadNVCUDAMgmt([]string{nvcudaLibPath}) + return cHandles + } + if cudartLibPath != "" { + cHandles.deviceCount, cHandles.cudart, _ = LoadCUDARTMgmt([]string{cudartLibPath}) + return cHandles } - slog.Debug("Detecting GPUs") - nvcudaLibPaths := FindGPULibs(nvcudaMgmtName, nvcudaMgmtPatterns) + slog.Debug("searching for GPU discovery libraries for NVIDIA") + var cudartMgmtPatterns []string + + // Aligned with driver, we can't carry as payloads + nvcudaMgmtPatterns := NvcudaGlobs + + if runtime.GOOS == "windows" { + localAppData := os.Getenv("LOCALAPPDATA") + cudartMgmtPatterns = []string{filepath.Join(localAppData, "Programs", "Ollama", CudartMgmtName)} + } + tmpDir, _ := PayloadsDir() + if tmpDir != "" { + // TODO - add "payloads" for subprocess + cudartMgmtPatterns = []string{filepath.Join(tmpDir, "cuda*", CudartMgmtName)} + } + cudartMgmtPatterns = append(cudartMgmtPatterns, CudartGlobs...) + + if len(NvmlGlobs) > 0 { + nvmlLibPaths := FindGPULibs(NvmlMgmtName, NvmlGlobs) + if len(nvmlLibPaths) > 0 { + nvml, libPath := LoadNVMLMgmt(nvmlLibPaths) + if nvml != nil { + slog.Debug("nvidia-ml loaded", "library", libPath) + cHandles.nvml = nvml + nvmlLibPath = libPath + } + } + } + + nvcudaLibPaths := FindGPULibs(NvcudaMgmtName, nvcudaMgmtPatterns) if len(nvcudaLibPaths) > 0 { deviceCount, nvcuda, libPath := LoadNVCUDAMgmt(nvcudaLibPaths) if nvcuda != nil { slog.Debug("detected GPUs", "count", deviceCount, "library", libPath) - gpuHandles.nvcuda = nvcuda - gpuHandles.deviceCount = deviceCount - return gpuHandles + cHandles.nvcuda = nvcuda + cHandles.deviceCount = deviceCount + nvcudaLibPath = libPath + return cHandles } } - cudartLibPaths := FindGPULibs(cudartMgmtName, cudartMgmtPatterns) + cudartLibPaths := FindGPULibs(CudartMgmtName, cudartMgmtPatterns) if len(cudartLibPaths) > 0 { deviceCount, cudart, libPath := LoadCUDARTMgmt(cudartLibPaths) if cudart != nil { slog.Debug("detected GPUs", "library", libPath, "count", deviceCount) - gpuHandles.cudart = cudart - gpuHandles.deviceCount = deviceCount - return gpuHandles + cHandles.cudart = cudart + cHandles.deviceCount = deviceCount + cudartLibPath = libPath + return cHandles } } - vulkanLibPaths := FindGPULibs(vulkanMgmtName, vulkanMgmtPatterns) + return cHandles +} - var libcapLibPaths []string - if runtime.GOOS == "linux" { - libcapLibPaths = FindGPULibs(libcapMgmtName, libcapMgmtPatterns) +// Note: gpuMutex must already be held +func initOneAPIHandles() *oneapiHandles { + oHandles := &oneapiHandles{} + + // Short Circuit if we already know which library to use + if oneapiLibPath != "" { + oHandles.deviceCount, oHandles.oneapi, _ = LoadOneapiMgmt([]string{oneapiLibPath}) + return oHandles + } + + oneapiLibPaths := FindGPULibs(OneapiMgmtName, OneapiGlobs) + if len(oneapiLibPaths) > 0 { + oHandles.deviceCount, oHandles.oneapi, oneapiLibPath = LoadOneapiMgmt(oneapiLibPaths) + } + + return oHandles +} + +// Note: gpuMutex must already be held +func initVulkanHandles() *vulkanHandles { + vHandles := &vulkanHandles{} + + // Short Circuit if we already know which library to use + if vulkanLibPath != "" && libcapLibPath != "" { + vHandles.deviceCount, vHandles.vulkan, _, _ = LoadVulkanMgmt([]string{vulkanLibPath}, []string{libcapLibPath}) + return vHandles + } + + vulkanPaths := FindGPULibs(VulkanMgmtName, VulkanGlobs) + libcapPaths := FindLibCapLibs() + + if len(vulkanPaths) > 0 && len(libcapPaths) > 0 { + vHandles.deviceCount, vHandles.vulkan, vulkanLibPath, libcapLibPath = LoadVulkanMgmt(vulkanPaths, libcapPaths) + } + + return vHandles +} + +func GetCPUInfo() GpuInfoList { + gpuMutex.Lock() + if !bootstrapped { + gpuMutex.Unlock() + GetGPUInfo() } else { - libcapLibPaths = []string{""} + gpuMutex.Unlock() } - - if len(vulkanLibPaths) > 0 && len(libcapLibPaths) > 0 { - deviceCount, vulkan, vkLibPath, capLibPath := LoadVulkanMgmt(vulkanLibPaths, libcapLibPaths) - if vulkan != nil { - slog.Debug("detected GPUs", "library", vkLibPath, capLibPath, "count", deviceCount) - gpuHandles.vulkan = vulkan - gpuHandles.deviceCount = deviceCount - return gpuHandles - } - } - - return gpuHandles + return GpuInfoList{cpus[0].GpuInfo} } func GetGPUInfo() GpuInfoList { @@ -200,141 +206,300 @@ func GetGPUInfo() GpuInfoList { // GPUs so we can report warnings if we see Nvidia/AMD but fail to load the libraries gpuMutex.Lock() defer gpuMutex.Unlock() - - gpuHandles := initGPUHandles() + needRefresh := true + var cHandles *cudaHandles + var oHandles *oneapiHandles + var vHandles *vulkanHandles defer func() { - if gpuHandles.cudart != nil { - C.cudart_release(*gpuHandles.cudart) + if cHandles != nil { + if cHandles.cudart != nil { + C.cudart_release(*cHandles.cudart) + } + if cHandles.nvcuda != nil { + C.nvcuda_release(*cHandles.nvcuda) + } + if cHandles.nvml != nil { + C.nvml_release(*cHandles.nvml) + } } - if gpuHandles.nvcuda != nil { - C.nvcuda_release(*gpuHandles.nvcuda) - } - if gpuHandles.vulkan != nil { - C.vk_release(*gpuHandles.vulkan) + if oHandles != nil { + if oHandles.oneapi != nil { + // TODO - is this needed? + C.oneapi_release(*oHandles.oneapi) + } } }() - // All our GPU builds on x86 have AVX enabled, so fallback to CPU if we don't detect at least AVX - cpuVariant := GetCPUVariant() - if cpuVariant == "" && runtime.GOARCH == "amd64" { - slog.Warn("CPU does not have AVX or AVX2, disabling GPU support.") - } + if !bootstrapped { + slog.Debug("Detecting GPUs") + needRefresh = false + cpuCapability = GetCPUCapability() + var memInfo C.mem_info_t - // On windows we bundle the nvidia library one level above the runner dir - depPath := "" - if runtime.GOOS == "windows" && envconfig.RunnersDir != "" { - depPath = filepath.Dir(envconfig.RunnersDir) - } - - var memInfo C.mem_info_t - resp := []GpuInfo{} - - // NVIDIA and Vulkan first - for i := range gpuHandles.deviceCount { - // TODO once we support CPU compilation variants of GPU libraries refine this... - if cpuVariant == "" && runtime.GOARCH == "amd64" { - continue + mem, err := GetCPUMem() + if err != nil { + slog.Warn("error looking up system memory", "error", err) } - if gpuHandles.cudart != nil || gpuHandles.nvcuda != nil { - gpuInfo := GpuInfo{ - Library: "cuda", + cpus = []CPUInfo{CPUInfo{ + GpuInfo: GpuInfo{ + memInfo: mem, + Library: "cpu", + Variant: cpuCapability, + ID: "0", + }, + }} + + // Fallback to CPU mode if we're lacking required vector extensions on x86 + if cpuCapability < GPURunnerCPUCapability && runtime.GOARCH == "amd64" { + slog.Warn("CPU does not have minimum vector extensions, GPU inference disabled", "required", GPURunnerCPUCapability, "detected", cpuCapability) + bootstrapped = true + // No need to do any GPU discovery, since we can't run on them + return GpuInfoList{cpus[0].GpuInfo} + } + + // On windows we bundle the nvidia library one level above the runner dir + depPath := "" + if runtime.GOOS == "windows" && envconfig.RunnersDir != "" { + depPath = filepath.Dir(envconfig.RunnersDir) + } + + // Load ALL libraries + cHandles = initCudaHandles() + + // NVIDIA + for i := range cHandles.deviceCount { + if cHandles.cudart != nil || cHandles.nvcuda != nil { + gpuInfo := CudaGPUInfo{ + GpuInfo: GpuInfo{ + Library: "cuda", + }, + index: i, + } + var driverMajor int + var driverMinor int + if cHandles.cudart != nil { + C.cudart_bootstrap(*cHandles.cudart, C.int(i), &memInfo) + } else { + C.nvcuda_bootstrap(*cHandles.nvcuda, C.int(i), &memInfo) + driverMajor = int(cHandles.nvcuda.driver_major) + driverMinor = int(cHandles.nvcuda.driver_minor) + } + if memInfo.err != nil { + slog.Info("error looking up nvidia GPU memory", "error", C.GoString(memInfo.err)) + C.free(unsafe.Pointer(memInfo.err)) + continue + } + if memInfo.major < CudaComputeMin[0] || (memInfo.major == CudaComputeMin[0] && memInfo.minor < CudaComputeMin[1]) { + slog.Info(fmt.Sprintf("[%d] CUDA GPU is too old. Compute Capability detected: %d.%d", i, memInfo.major, memInfo.minor)) + continue + } + gpuInfo.TotalMemory = uint64(memInfo.total) + gpuInfo.FreeMemory = uint64(memInfo.free) + gpuInfo.ID = C.GoString(&memInfo.gpu_id[0]) + gpuInfo.Compute = fmt.Sprintf("%d.%d", memInfo.major, memInfo.minor) + gpuInfo.MinimumMemory = cudaMinimumMemory + gpuInfo.DependencyPath = depPath + gpuInfo.Name = C.GoString(&memInfo.gpu_name[0]) + gpuInfo.DriverMajor = driverMajor + gpuInfo.DriverMinor = driverMinor + + // TODO potentially sort on our own algorithm instead of what the underlying GPU library does... + cudaGPUs = append(cudaGPUs, gpuInfo) } - var driverMajor int - var driverMinor int - if gpuHandles.cudart != nil { - C.cudart_check_vram(*gpuHandles.cudart, C.int(i), &memInfo) + } + + // Intel + oHandles = initOneAPIHandles() + for d := 0; oHandles.oneapi != nil && d < int(oHandles.oneapi.num_drivers); d++ { + if oHandles.oneapi == nil { + // shouldn't happen + slog.Warn("nil oneapi handle with driver count", "count", int(oHandles.oneapi.num_drivers)) + continue + } + devCount := C.oneapi_get_device_count(*oHandles.oneapi, C.int(d)) + for i := range devCount { + gpuInfo := OneapiGPUInfo{ + GpuInfo: GpuInfo{ + Library: "oneapi", + }, + driverIndex: d, + gpuIndex: int(i), + } + // TODO - split bootstrapping from updating free memory + C.oneapi_check_vram(*oHandles.oneapi, C.int(d), i, &memInfo) + // TODO - convert this to MinimumMemory based on testing... + var totalFreeMem float64 = float64(memInfo.free) * 0.95 // work-around: leave some reserve vram for mkl lib used in ggml-sycl backend. + memInfo.free = C.uint64_t(totalFreeMem) + gpuInfo.TotalMemory = uint64(memInfo.total) + gpuInfo.FreeMemory = uint64(memInfo.free) + gpuInfo.ID = C.GoString(&memInfo.gpu_id[0]) + gpuInfo.Name = C.GoString(&memInfo.gpu_name[0]) + // TODO dependency path? + oneapiGPUs = append(oneapiGPUs, gpuInfo) + } + } + + // Vulkan + vHandles = initVulkanHandles() + for i := range vHandles.deviceCount { + if vHandles.vulkan != nil { + gpuInfo := VulkanGPUInfo{ + GpuInfo: GpuInfo{ + Library: "vulkan", + }, + index: i, + } + + C.vk_check_vram(*vHandles.vulkan, C.int(i), &memInfo) + if memInfo.err != nil { + slog.Info("error looking up vulkan GPU memory", "error", C.GoString(memInfo.err)) + C.free(unsafe.Pointer(memInfo.err)) + continue + } + + gpuInfo.TotalMemory = uint64(memInfo.total) + gpuInfo.FreeMemory = uint64(memInfo.free) + gpuInfo.ID = C.GoString(&memInfo.gpu_id[0]) + gpuInfo.Compute = fmt.Sprintf("%d.%d", memInfo.major, memInfo.minor) + gpuInfo.MinimumMemory = 0 + gpuInfo.DependencyPath = depPath + gpuInfo.Name = C.GoString(&memInfo.gpu_name[0]) + gpuInfo.DriverMajor = int(memInfo.major) + gpuInfo.DriverMinor = int(memInfo.minor) + + // TODO potentially sort on our own algorithm instead of what the underlying GPU library does... + vulkanGPUs = append(vulkanGPUs, gpuInfo) + } + } + + rocmGPUs = AMDGetGPUInfo() + bootstrapped = true + } + + // For detected GPUs, load library if not loaded + + // Refresh free memory usage + if needRefresh { + mem, err := GetCPUMem() + if err != nil { + slog.Warn("error looking up system memory", "error", err) + } else { + slog.Debug("updating system memory data", + slog.Group( + "before", + "total", format.HumanBytes2(cpus[0].TotalMemory), + "free", format.HumanBytes2(cpus[0].FreeMemory), + ), + slog.Group( + "now", + "total", format.HumanBytes2(mem.TotalMemory), + "free", format.HumanBytes2(mem.FreeMemory), + ), + ) + cpus[0].FreeMemory = mem.FreeMemory + } + + var memInfo C.mem_info_t + if cHandles == nil && len(cudaGPUs) > 0 { + cHandles = initCudaHandles() + } + for i, gpu := range cudaGPUs { + if cHandles.nvml != nil { + C.nvml_get_free(*cHandles.nvml, C.int(gpu.index), &memInfo.free, &memInfo.total, &memInfo.used) + } else if cHandles.cudart != nil { + C.cudart_bootstrap(*cHandles.cudart, C.int(gpu.index), &memInfo) + } else if cHandles.nvcuda != nil { + C.nvcuda_get_free(*cHandles.nvcuda, C.int(gpu.index), &memInfo.free, &memInfo.total) + memInfo.used = memInfo.total - memInfo.free } else { - C.nvcuda_check_vram(*gpuHandles.nvcuda, C.int(i), &memInfo) - driverMajor = int(gpuHandles.nvcuda.driver_major) - driverMinor = int(gpuHandles.nvcuda.driver_minor) + // shouldn't happen + slog.Warn("no valid cuda library loaded to refresh vram usage") + break } if memInfo.err != nil { - slog.Info("error looking up nvidia GPU memory", "error", C.GoString(memInfo.err)) + slog.Warn("error looking up nvidia GPU memory", "error", C.GoString(memInfo.err)) C.free(unsafe.Pointer(memInfo.err)) continue } - if memInfo.major < CudaComputeMin[0] || (memInfo.major == CudaComputeMin[0] && memInfo.minor < CudaComputeMin[1]) { - slog.Info(fmt.Sprintf("[%d] CUDA GPU is too old. Compute Capability detected: %d.%d", i, memInfo.major, memInfo.minor)) + if memInfo.free == 0 { + slog.Warn("error looking up nvidia GPU memory") continue } - gpuInfo.TotalMemory = uint64(memInfo.total) - gpuInfo.FreeMemory = uint64(memInfo.free) - gpuInfo.ID = C.GoString(&memInfo.gpu_id[0]) - gpuInfo.Compute = fmt.Sprintf("%d.%d", memInfo.major, memInfo.minor) - gpuInfo.MinimumMemory = cudaMinimumMemory - gpuInfo.DependencyPath = depPath - gpuInfo.Name = C.GoString(&memInfo.gpu_name[0]) - gpuInfo.DriverMajor = driverMajor - gpuInfo.DriverMinor = driverMinor - - // TODO potentially sort on our own algorithm instead of what the underlying GPU library does... - resp = append(resp, gpuInfo) + slog.Debug("updating cuda memory data", + "gpu", gpu.ID, + "name", gpu.Name, + slog.Group( + "before", + "total", format.HumanBytes2(gpu.TotalMemory), + "free", format.HumanBytes2(gpu.FreeMemory), + ), + slog.Group( + "now", + "total", format.HumanBytes2(uint64(memInfo.total)), + "free", format.HumanBytes2(uint64(memInfo.free)), + "used", format.HumanBytes2(uint64(memInfo.used)), + ), + ) + cudaGPUs[i].FreeMemory = uint64(memInfo.free) } - if gpuHandles.vulkan != nil { - gpuInfo := GpuInfo{ - Library: "vulkan", - } - - C.vk_check_vram(*gpuHandles.vulkan, C.int(i), &memInfo) - if memInfo.err != nil { - slog.Info("error looking up vulkan GPU memory", "error", C.GoString(memInfo.err)) - C.free(unsafe.Pointer(memInfo.err)) + if oHandles == nil && len(oneapiGPUs) > 0 { + oHandles = initOneAPIHandles() + } + for i, gpu := range oneapiGPUs { + if oHandles.oneapi == nil { + // shouldn't happen + slog.Warn("nil oneapi handle with device count", "count", oHandles.deviceCount) continue } + C.oneapi_check_vram(*oHandles.oneapi, C.int(gpu.driverIndex), C.int(gpu.gpuIndex), &memInfo) + // TODO - convert this to MinimumMemory based on testing... + var totalFreeMem float64 = float64(memInfo.free) * 0.95 // work-around: leave some reserve vram for mkl lib used in ggml-sycl backend. + memInfo.free = C.uint64_t(totalFreeMem) + oneapiGPUs[i].FreeMemory = uint64(memInfo.free) + } - gpuInfo.TotalMemory = uint64(memInfo.total) - gpuInfo.FreeMemory = uint64(memInfo.free) - gpuInfo.ID = C.GoString(&memInfo.gpu_id[0]) - gpuInfo.Compute = fmt.Sprintf("%d.%d", memInfo.major, memInfo.minor) - gpuInfo.MinimumMemory = 0 - gpuInfo.DependencyPath = depPath - gpuInfo.Name = C.GoString(&memInfo.gpu_name[0]) - gpuInfo.DriverMajor = int(memInfo.major) - gpuInfo.DriverMinor = int(memInfo.minor) + if vHandles == nil && len(vulkanGPUs) > 0 { + vHandles = initVulkanHandles() + } + for i, gpu := range vulkanGPUs { + if vHandles.vulkan == nil { + // shouldn't happen + slog.Warn("nil vulkan handle with device count", "count", oHandles.deviceCount) + continue + } + C.vk_check_vram(*vHandles.vulkan, C.int(gpu.index), &memInfo) + // TODO - convert this to MinimumMemory based on testing... + var totalFreeMem float64 = float64(memInfo.free) * 0.95 // work-around: leave some reserve vram for mkl lib used in ggml-sycl backend. + memInfo.free = C.uint64_t(totalFreeMem) + vulkanGPUs[i].FreeMemory = uint64(memInfo.free) + } - // TODO potentially sort on our own algorithm instead of what the underlying GPU library does... - resp = append(resp, gpuInfo) + err = RocmGPUInfoList(rocmGPUs).RefreshFreeMemory() + if err != nil { + slog.Debug("problem refreshing ROCm free memory", "error", err) } } - // Then AMD - resp = append(resp, AMDGetGPUInfo()...) - + resp := []GpuInfo{} + for _, gpu := range cudaGPUs { + resp = append(resp, gpu.GpuInfo) + } + for _, gpu := range rocmGPUs { + resp = append(resp, gpu.GpuInfo) + } + for _, gpu := range oneapiGPUs { + resp = append(resp, gpu.GpuInfo) + } + for _, gpu := range vulkanGPUs { + resp = append(resp, gpu.GpuInfo) + } if len(resp) == 0 { - C.cpu_check_ram(&memInfo) - if memInfo.err != nil { - slog.Info("error looking up CPU memory", "error", C.GoString(memInfo.err)) - C.free(unsafe.Pointer(memInfo.err)) - return resp - } - gpuInfo := GpuInfo{ - Library: "cpu", - Variant: cpuVariant, - } - gpuInfo.TotalMemory = uint64(memInfo.total) - gpuInfo.FreeMemory = uint64(memInfo.free) - gpuInfo.ID = C.GoString(&memInfo.gpu_id[0]) - - resp = append(resp, gpuInfo) + resp = append(resp, cpus[0].GpuInfo) } - return resp } -func GetCPUMem() (memInfo, error) { - var ret memInfo - var info C.mem_info_t - C.cpu_check_ram(&info) - if info.err != nil { - defer C.free(unsafe.Pointer(info.err)) - return ret, fmt.Errorf(C.GoString(info.err)) - } - ret.FreeMemory = uint64(info.free) - ret.TotalMemory = uint64(info.total) - return ret, nil -} - func FindGPULibs(baseLibName string, defaultPatterns []string) []string { // Multiple GPU libraries may exist, and some may not work, so keep trying until we exhaust them var ldPaths []string @@ -431,8 +596,26 @@ func LoadNVCUDAMgmt(nvcudaLibPaths []string) (int, *C.nvcuda_handle_t, string) { return 0, nil, "" } +func LoadNVMLMgmt(nvmlLibPaths []string) (*C.nvml_handle_t, string) { + var resp C.nvml_init_resp_t + resp.ch.verbose = getVerboseState() + for _, libPath := range nvmlLibPaths { + lib := C.CString(libPath) + defer C.free(unsafe.Pointer(lib)) + C.nvml_init(lib, &resp) + if resp.err != nil { + slog.Info(fmt.Sprintf("Unable to load NVML management library %s: %s", libPath, C.GoString(resp.err))) + C.free(unsafe.Pointer(resp.err)) + } else { + return &resp.ch, libPath + } + } + return nil, "" +} + func LoadOneapiMgmt(oneapiLibPaths []string) (int, *C.oneapi_handle_t, string) { var resp C.oneapi_init_resp_t + num_devices := 0 resp.oh.verbose = getVerboseState() for _, libPath := range oneapiLibPaths { lib := C.CString(libPath) @@ -442,7 +625,10 @@ func LoadOneapiMgmt(oneapiLibPaths []string) (int, *C.oneapi_handle_t, string) { slog.Debug("Unable to load oneAPI management library", "library", libPath, "error", C.GoString(resp.err)) C.free(unsafe.Pointer(resp.err)) } else { - return int(resp.num_devices), &resp.oh, libPath + for i := range resp.oh.num_drivers { + num_devices += int(C.oneapi_get_device_count(resp.oh, C.int(i))) + } + return num_devices, &resp.oh, libPath } } return 0, nil, "" diff --git a/gpu/gpu_linux.go b/gpu/gpu_linux.go index a099bf822..2e723c4da 100644 --- a/gpu/gpu_linux.go +++ b/gpu/gpu_linux.go @@ -43,10 +43,26 @@ var OneapiGlobs = []string{ "/usr/lib*/libze_intel_gpu.so*", } +var VulkanGlobs = []string{ + "/usr/lib/x86_64-linux-gnu/libvulkan.so*", + "/usr/lib*/libvulkan.so*", +} + +var capLinuxGlobs = []string{ + "/usr/lib/x86_64-linux-gnu/libcap.so*", + "/usr/lib*/libcap.so*", +} + var CudartMgmtName = "libcudart.so*" var NvcudaMgmtName = "libcuda.so*" var NvmlMgmtName = "" // not currently wired on linux var OneapiMgmtName = "libze_intel_gpu.so" +var VulkanMgmtName = "libvulkan.so*" +var libcapMgmtName = "libcap.so*" + +func FindLibCapLibs() []string { + return FindGPULibs(libcapMgmtName, capLinuxGlobs) +} func GetCPUMem() (memInfo, error) { var mem memInfo diff --git a/gpu/gpu_windows.go b/gpu/gpu_windows.go index f8c2e76fe..328477440 100644 --- a/gpu/gpu_windows.go +++ b/gpu/gpu_windows.go @@ -45,6 +45,10 @@ var NvcudaMgmtName = "nvcuda.dll" var NvmlMgmtName = "nvml.dll" var OneapiMgmtName = "ze_intel_gpu64.dll" +func FindLibCapLibs() []string { + return []string{""} +} + func GetCPUMem() (memInfo, error) { memStatus := MEMORYSTATUSEX{length: sizeofMemoryStatusEx} r1, _, err := globalMemoryStatusExProc.Call(uintptr(unsafe.Pointer(&memStatus))) diff --git a/gpu/types.go b/gpu/types.go index 47355959c..b451c0f38 100644 --- a/gpu/types.go +++ b/gpu/types.go @@ -62,6 +62,13 @@ type OneapiGPUInfo struct { } type OneapiGPUInfoList []OneapiGPUInfo +type VulkanGPUInfo struct { + GpuInfo + index int +} + +type VulkanGPUInfoList []VulkanGPUInfo + type GpuInfoList []GpuInfo // Split up the set of gpu info's by Library and variant From 38466f18213ae6f6879e5639d315e8c8e377b602 Mon Sep 17 00:00:00 2001 From: pufferffish Date: Sat, 15 Jun 2024 12:06:43 +0100 Subject: [PATCH 10/29] fix build --- gpu/gpu.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gpu/gpu.go b/gpu/gpu.go index 11c72e151..4eb82f2cf 100644 --- a/gpu/gpu.go +++ b/gpu/gpu.go @@ -37,7 +37,7 @@ type oneapiHandles struct { } type vulkanHandles struct { - vulkan *C.vulkan_handle_t + vulkan *C.vk_handle_t deviceCount int } From e3f9ca4009afe2620b27b61b3e0f37053b1d4354 Mon Sep 17 00:00:00 2001 From: KOISHI KOMEIJI FROM TOUHOU 11 Date: Sat, 15 Jun 2024 20:13:15 +0800 Subject: [PATCH 11/29] fix check_perfmon len --- gpu/gpu_info_vulkan.c | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gpu/gpu_info_vulkan.c b/gpu/gpu_info_vulkan.c index fbe7a5885..17ee43003 100644 --- a/gpu/gpu_info_vulkan.c +++ b/gpu/gpu_info_vulkan.c @@ -5,7 +5,7 @@ int check_perfmon(vk_handle_t* rh) { #ifdef __linux__ cap_t caps; - const cap_value_t cap_list[2] = {CAP_PERFMON}; + const cap_value_t cap_list[1] = {CAP_PERFMON}; if ((*rh->cap_get_bound)(CAP_SETFCAP) < 0) return -1; @@ -14,7 +14,7 @@ int check_perfmon(vk_handle_t* rh) { if (caps == NULL) return -1; - if ((*rh->cap_set_flag)(caps, CAP_EFFECTIVE, 2, cap_list, CAP_SET) == -1) + if ((*rh->cap_set_flag)(caps, CAP_EFFECTIVE, 1, cap_list, CAP_SET) == -1) return -1; if ((*rh->cap_set_proc)(caps) == -1) @@ -221,4 +221,4 @@ void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp) { void vk_release(vk_handle_t rh) { (*rh.vkDestroyInstance)(rh.vk, NULL); -} \ No newline at end of file +} From b958cd2848773e1e37fe0cd1e000aa0ee65f5fff Mon Sep 17 00:00:00 2001 From: DSLstandard Date: Sat, 15 Jun 2024 20:19:19 +0800 Subject: [PATCH 12/29] remove cap_get_bound check --- gpu/gpu_info_vulkan.c | 3 --- 1 file changed, 3 deletions(-) diff --git a/gpu/gpu_info_vulkan.c b/gpu/gpu_info_vulkan.c index 17ee43003..c4cdaa543 100644 --- a/gpu/gpu_info_vulkan.c +++ b/gpu/gpu_info_vulkan.c @@ -7,9 +7,6 @@ int check_perfmon(vk_handle_t* rh) { cap_t caps; const cap_value_t cap_list[1] = {CAP_PERFMON}; - if ((*rh->cap_get_bound)(CAP_SETFCAP) < 0) - return -1; - caps = (*rh->cap_get_proc)(); if (caps == NULL) return -1; From b6554e9b8c6502e2dcebba17bab75bd5235adfff Mon Sep 17 00:00:00 2001 From: pufferffish Date: Sat, 15 Jun 2024 21:11:07 +0100 Subject: [PATCH 13/29] fix vulkan handle releasing --- gpu/gpu.go | 5 +++++ gpu/gpu_info_vulkan.c | 9 ++++++++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/gpu/gpu.go b/gpu/gpu.go index 4eb82f2cf..6cebbd2b9 100644 --- a/gpu/gpu.go +++ b/gpu/gpu.go @@ -228,6 +228,11 @@ func GetGPUInfo() GpuInfoList { C.oneapi_release(*oHandles.oneapi) } } + if vHandles != nil { + if vHandles.vulkan != nil { + C.vk_release(*vHandles.vulkan) + } + } }() if !bootstrapped { diff --git a/gpu/gpu_info_vulkan.c b/gpu/gpu_info_vulkan.c index c4cdaa543..e868dcc1b 100644 --- a/gpu/gpu_info_vulkan.c +++ b/gpu/gpu_info_vulkan.c @@ -213,9 +213,16 @@ void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp) { resp->major = VK_API_VERSION_MAJOR(properties.apiVersion); resp->minor = VK_API_VERSION_MINOR(properties.apiVersion); resp->patch = VK_API_VERSION_PATCH(properties.apiVersion); - } void vk_release(vk_handle_t rh) { + LOG(rh.verbose, "releasing vulkan library\n"); (*rh.vkDestroyInstance)(rh.vk, NULL); + UNLOAD_LIBRARY(rh.vk_handle); + rh.vk_handle = NULL; +#ifdef __linux__ + LOG(rh.verbose, "releasing libcap library\n"); + UNLOAD_LIBRARY(rh.cap_handle); + rh.cap_handle = NULL; +#endif } From ace3d104683748e627b284c21bdfe387536f3b59 Mon Sep 17 00:00:00 2001 From: pufferffish Date: Mon, 23 Sep 2024 18:38:42 +0800 Subject: [PATCH 14/29] fix build on federa 40 --- llm/generate/gen_linux.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/llm/generate/gen_linux.sh b/llm/generate/gen_linux.sh index 7e5f531e1..79c449018 100755 --- a/llm/generate/gen_linux.sh +++ b/llm/generate/gen_linux.sh @@ -234,8 +234,8 @@ if [ -z "${OLLAMA_SKIP_VULKAN_GENERATE}" -a -d "${VULKAN_ROOT}" ] && [ -z "${OLL for dep in $(ldd "${BUILD_DIR}/bin/ollama_llama_server" | grep "=>" | cut -f2 -d= | cut -f2 -d' ' | grep -e vulkan -e cap); do cp "${dep}" "${BUILD_DIR}/bin/" done - cp "${VULKAN_ROOT}/libvulkan.so" "${BUILD_DIR}/bin/" - cp "${CAP_ROOT}/libcap.so" "${BUILD_DIR}/bin/" + cp "${VULKAN_ROOT}/libvulkan.so*" "${BUILD_DIR}/bin/" + cp "${CAP_ROOT}/libcap.so*" "${BUILD_DIR}/bin/" compress fi From e61c32943556142a646e5812654768d9c2ba29c8 Mon Sep 17 00:00:00 2001 From: pufferffish Date: Mon, 23 Sep 2024 18:43:16 +0800 Subject: [PATCH 15/29] fix vulkan on windows --- gpu/gpu_windows.go | 5 +++++ llm/generate/gen_windows.ps1 | 16 ++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/gpu/gpu_windows.go b/gpu/gpu_windows.go index 5491da963..f9afb6ed8 100644 --- a/gpu/gpu_windows.go +++ b/gpu/gpu_windows.go @@ -40,11 +40,16 @@ var OneapiGlobs = []string{ "c:\\Windows\\System32\\DriverStore\\FileRepository\\*\\ze_intel_gpu64.dll", } +var VulkanGlobs = []string{ + "c:\\Windows\\System32\\vulkan-1.dll", +} + var ( CudartMgmtName = "cudart64_*.dll" NvcudaMgmtName = "nvcuda.dll" NvmlMgmtName = "nvml.dll" OneapiMgmtName = "ze_intel_gpu64.dll" + VulkanMgmtName = "vulkan-1.dll" ) func FindLibCapLibs() []string { diff --git a/llm/generate/gen_windows.ps1 b/llm/generate/gen_windows.ps1 index 29ff5ff62..bb92c5121 100644 --- a/llm/generate/gen_windows.ps1 +++ b/llm/generate/gen_windows.ps1 @@ -412,6 +412,21 @@ function build_rocm() { } } +function build_vulkan() { + if (-not "${env:OLLAMA_SKIP_VULKAN_GENERATE}") { + init_vars + $script:buildDir="../build/windows/${script:ARCH}/vulkan" + $script:distDir="$script:DIST_BASE\vulkan" + $script:cmakeDefs += @("-A", "x64", "-DLLAMA_VULKAN=1") + write-host "Building Vulkan" + build + sign + install + } else { + write-host "Skipping Vulkan generation step" + } +} + init_vars if ($($args.count) -eq 0) { git_module_setup @@ -426,6 +441,7 @@ if ($($args.count) -eq 0) { build_cuda build_oneapi build_rocm + build_vulkan } cleanup From 4b74cee096f9bed2d5b64a983575c9d9d4c6fe7d Mon Sep 17 00:00:00 2001 From: yeongbba Date: Sun, 19 Jan 2025 01:30:34 +0900 Subject: [PATCH 16/29] making amdgpu work on arm achitecutre with vulkan --- envconfig/config.go | 2 ++ gpu/gpu.go | 2 +- gpu/gpu_linux.go | 4 ++-- llm/generate/gen_common.sh | 4 ++-- llm/generate/gen_linux.sh | 12 +++++++----- llm/generate/gen_windows.ps1 | 2 +- 6 files changed, 15 insertions(+), 11 deletions(-) diff --git a/envconfig/config.go b/envconfig/config.go index 9c1490a93..239c49fe6 100644 --- a/envconfig/config.go +++ b/envconfig/config.go @@ -175,6 +175,7 @@ var ( CudaVisibleDevices = String("CUDA_VISIBLE_DEVICES") HipVisibleDevices = String("HIP_VISIBLE_DEVICES") RocrVisibleDevices = String("ROCR_VISIBLE_DEVICES") + VkVisibleDevices = String("GGML_VK_VISIBLE_DEVICES") GpuDeviceOrdinal = String("GPU_DEVICE_ORDINAL") HsaOverrideGfxVersion = String("HSA_OVERRIDE_GFX_VERSION") ) @@ -263,6 +264,7 @@ func AsMap() map[string]EnvVar { ret["CUDA_VISIBLE_DEVICES"] = EnvVar{"CUDA_VISIBLE_DEVICES", CudaVisibleDevices(), "Set which NVIDIA devices are visible"} ret["HIP_VISIBLE_DEVICES"] = EnvVar{"HIP_VISIBLE_DEVICES", HipVisibleDevices(), "Set which AMD devices are visible"} ret["ROCR_VISIBLE_DEVICES"] = EnvVar{"ROCR_VISIBLE_DEVICES", RocrVisibleDevices(), "Set which AMD devices are visible"} + ret["GGML_VK_VISIBLE_DEVICES"] = EnvVar{"GGML_VK_VISIBLE_DEVICES", VkVisibleDevices(), "Set which VK AMD devices are visible"} ret["GPU_DEVICE_ORDINAL"] = EnvVar{"GPU_DEVICE_ORDINAL", GpuDeviceOrdinal(), "Set which AMD devices are visible"} ret["HSA_OVERRIDE_GFX_VERSION"] = EnvVar{"HSA_OVERRIDE_GFX_VERSION", HsaOverrideGfxVersion(), "Override the gfx used for all detected AMD GPUs"} ret["OLLAMA_INTEL_GPU"] = EnvVar{"OLLAMA_INTEL_GPU", IntelGPU(), "Enable experimental Intel GPU detection"} diff --git a/gpu/gpu.go b/gpu/gpu.go index 69279867e..1bd337f19 100644 --- a/gpu/gpu.go +++ b/gpu/gpu.go @@ -410,7 +410,7 @@ func GetGPUInfo() GpuInfoList { rocmGPUs = AMDGetGPUInfo() bootstrapped = true - if len(cudaGPUs) == 0 && len(rocmGPUs) == 0 && len(oneapiGPUs) == 0 { + if len(cudaGPUs) == 0 && len(rocmGPUs) == 0 && len(oneapiGPUs) == 0 && len(vulkanGPUs) == 0 { slog.Info("no compatible GPUs were discovered") } } diff --git a/gpu/gpu_linux.go b/gpu/gpu_linux.go index 76df63268..d6f882efb 100644 --- a/gpu/gpu_linux.go +++ b/gpu/gpu_linux.go @@ -53,12 +53,12 @@ var ( ) var VulkanGlobs = []string{ - "/usr/lib/x86_64-linux-gnu/libvulkan.so*", + "/usr/lib/aarch64-linux-gnu/libvulkan.so*", "/usr/lib*/libvulkan.so*", } var capLinuxGlobs = []string{ - "/usr/lib/x86_64-linux-gnu/libcap.so*", + "/usr/lib/aarch64-linux-gnu/libcap.so*", "/usr/lib*/libcap.so*", } diff --git a/llm/generate/gen_common.sh b/llm/generate/gen_common.sh index 3825c155a..2b01e149c 100644 --- a/llm/generate/gen_common.sh +++ b/llm/generate/gen_common.sh @@ -30,7 +30,7 @@ init_vars() { WHOLE_ARCHIVE="-Wl,-force_load" NO_WHOLE_ARCHIVE="" GCC_ARCH="-arch ${ARCH}" - DIST_BASE=../../dist/darwin-${GOARCH}/ + DIST_BASE=../../dist/darwin-${GOARCH} PAYLOAD_BASE=../../build/darwin/${GOARCH} ;; "Linux") @@ -40,7 +40,7 @@ init_vars() { # Cross compiling not supported on linux - Use docker GCC_ARCH="" - DIST_BASE=../../dist/linux-${GOARCH}/ + DIST_BASE=../../dist/linux-${GOARCH} PAYLOAD_BASE=../../build/linux/${GOARCH} ;; *) diff --git a/llm/generate/gen_linux.sh b/llm/generate/gen_linux.sh index 79c449018..17981d543 100755 --- a/llm/generate/gen_linux.sh +++ b/llm/generate/gen_linux.sh @@ -224,9 +224,9 @@ fi if [ -z "${OLLAMA_SKIP_VULKAN_GENERATE}" -a -d "${VULKAN_ROOT}" ] && [ -z "${OLLAMA_SKIP_VULKAN_GENERATE}" -a -d "${CAP_ROOT}" ]; then echo "Vulkan and capabilities libraries detected - building dynamic Vulkan library" init_vars - - CMAKE_DEFS="${COMMON_CMAKE_DEFS} ${CMAKE_DEFS} -DLLAMA_VULKAN=1" - BUILD_DIR="../build/linux/${ARCH}/vulkan" + RUNNER=vulkan + CMAKE_DEFS="-DCMAKE_SKIP_RPATH=on -DBUILD_SHARED_LIBS=on -DCMAKE_POSITION_INDEPENDENT_CODE=on -DGGML_NATIVE=off -DGGML_AVX=off -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off -DGGML_OPENMP=off" + BUILD_DIR="../build/linux/${ARCH}/${RUNNER}" EXTRA_LIBS="-L${VULKAN_ROOT} -L${CAP_ROOT} -lvulkan -lcap" build @@ -234,8 +234,10 @@ if [ -z "${OLLAMA_SKIP_VULKAN_GENERATE}" -a -d "${VULKAN_ROOT}" ] && [ -z "${OLL for dep in $(ldd "${BUILD_DIR}/bin/ollama_llama_server" | grep "=>" | cut -f2 -d= | cut -f2 -d' ' | grep -e vulkan -e cap); do cp "${dep}" "${BUILD_DIR}/bin/" done - cp "${VULKAN_ROOT}/libvulkan.so*" "${BUILD_DIR}/bin/" - cp "${CAP_ROOT}/libcap.so*" "${BUILD_DIR}/bin/" + cp ${VULKAN_ROOT}/libvulkan.so* "${BUILD_DIR}/bin/" + cp ${CAP_ROOT}/libcap.so* "${BUILD_DIR}/bin/" + install + dist compress fi diff --git a/llm/generate/gen_windows.ps1 b/llm/generate/gen_windows.ps1 index bb92c5121..e8bba29a5 100644 --- a/llm/generate/gen_windows.ps1 +++ b/llm/generate/gen_windows.ps1 @@ -417,7 +417,7 @@ function build_vulkan() { init_vars $script:buildDir="../build/windows/${script:ARCH}/vulkan" $script:distDir="$script:DIST_BASE\vulkan" - $script:cmakeDefs += @("-A", "x64", "-DLLAMA_VULKAN=1") + $script:cmakeDefs += @("-A", "x64", "-DDGGML_VULKAN=1") write-host "Building Vulkan" build sign From 6d7579b567cdb03154521e43103e9113a8cbd336 Mon Sep 17 00:00:00 2001 From: yeongbba Date: Sun, 19 Jan 2025 12:41:08 +0900 Subject: [PATCH 17/29] add x86_64 lines in VulkanGlobs and capLinuxGlobs --- gpu/gpu_linux.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/gpu/gpu_linux.go b/gpu/gpu_linux.go index d6f882efb..1251c6e8e 100644 --- a/gpu/gpu_linux.go +++ b/gpu/gpu_linux.go @@ -53,11 +53,13 @@ var ( ) var VulkanGlobs = []string{ + "/usr/lib/x86_64-linux-gnu/libvulkan.so*", "/usr/lib/aarch64-linux-gnu/libvulkan.so*", "/usr/lib*/libvulkan.so*", } var capLinuxGlobs = []string{ + "/usr/lib/x86_64-linux-gnu/libvulkan.so*", "/usr/lib/aarch64-linux-gnu/libcap.so*", "/usr/lib*/libcap.so*", } From 2bf59a512b938739f3bb8c3cdd034d817692ed45 Mon Sep 17 00:00:00 2001 From: yeongbba Date: Sun, 19 Jan 2025 12:51:10 +0900 Subject: [PATCH 18/29] add aarch64 lines in vulkanGlobs and capLinuxGlobs --- discover/gpu_linux.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/discover/gpu_linux.go b/discover/gpu_linux.go index 840ea435a..a2c1d8715 100644 --- a/discover/gpu_linux.go +++ b/discover/gpu_linux.go @@ -58,11 +58,13 @@ var ( var VulkanGlobs = []string{ "/usr/lib/x86_64-linux-gnu/libvulkan.so*", + "/usr/lib/aarch64-linux-gnu/libvulkan.so*", "/usr/lib*/libvulkan.so*", } var capLinuxGlobs = []string{ "/usr/lib/x86_64-linux-gnu/libcap.so*", + "/usr/lib/aarch64-linux-gnu/libvulkan.so*", "/usr/lib*/libcap.so*", } From 0d277d32db8a56aa42e2fdeea2c81202a364c8da Mon Sep 17 00:00:00 2001 From: tomaThomas Date: Sat, 25 Jan 2025 11:23:25 +0100 Subject: [PATCH 19/29] Fix variable name --- discover/gpu.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/discover/gpu.go b/discover/gpu.go index 15d0b99a4..1b6a3f075 100644 --- a/discover/gpu.go +++ b/discover/gpu.go @@ -440,7 +440,7 @@ func GetGPUInfo() GpuInfoList { gpuInfo.ID = C.GoString(&memInfo.gpu_id[0]) gpuInfo.Compute = fmt.Sprintf("%d.%d", memInfo.major, memInfo.minor) gpuInfo.MinimumMemory = 0 - gpuInfo.DependencyPath = depPath + gpuInfo.DependencyPath = depPaths gpuInfo.Name = C.GoString(&memInfo.gpu_name[0]) gpuInfo.DriverMajor = int(memInfo.major) gpuInfo.DriverMinor = int(memInfo.minor) From 2d443b3dd660a1fd2760d64538512df93648b4bb Mon Sep 17 00:00:00 2001 From: pufferffish Date: Mon, 3 Feb 2025 14:46:59 +0000 Subject: [PATCH 20/29] Add vulkan build patch from @jmorganca --- CMakeLists.txt | 13 +++++++++++++ ml/backend/ggml/ggml/.rsync-filter | 3 +++ 2 files changed, 16 insertions(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index 19d9bd8f9..05f8e2c47 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -110,3 +110,16 @@ if(CMAKE_HIP_COMPILER) endforeach() endif() endif() + +find_package(Vulkan) +if(Vulkan_FOUND) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/ml/backend/ggml/ggml/src/ggml-vulkan) + set(OLLAMA_VULKAN_INSTALL_DIR ${OLLAMA_INSTALL_DIR}/vulkan) + install(TARGETS ggml-vulkan + RUNTIME_DEPENDENCIES + PRE_INCLUDE_REGEXES vulkan + PRE_EXCLUDE_REGEXES ".*" + RUNTIME DESTINATION ${OLLAMA_VULKAN_INSTALL_DIR} COMPONENT Vulkan + LIBRARY DESTINATION ${OLLAMA_VULKAN_INSTALL_DIR} COMPONENT Vulkan + ) +endif() diff --git a/ml/backend/ggml/ggml/.rsync-filter b/ml/backend/ggml/ggml/.rsync-filter index c5acbe490..09d67f270 100644 --- a/ml/backend/ggml/ggml/.rsync-filter +++ b/ml/backend/ggml/ggml/.rsync-filter @@ -12,6 +12,8 @@ include src/ggml-cuda/ include src/ggml-cuda/template-instances/ include src/ggml-hip/ include src/ggml-metal/ +include src/ggml-vulkan/ +include src/ggml-vulkan/vulkan-shaders include *.c include *.h include *.cpp @@ -19,4 +21,5 @@ include *.cu include *.cuh include *.m include *.metal +include *.comp exclude * From 449e5c07aeb9c034f530869a55f0fe3f44ee88dc Mon Sep 17 00:00:00 2001 From: Antoine Viallon Date: Tue, 4 Feb 2025 11:51:17 +0100 Subject: [PATCH 21/29] Sync vendored ggml to add Vulkan support --- Makefile.sync | 2 +- .../ggml/ggml/src/ggml-vulkan/CMakeLists.txt | 92 + .../ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp | 8745 +++++++++++++++++ .../ggml-vulkan/vulkan-shaders/CMakeLists.txt | 9 + .../src/ggml-vulkan/vulkan-shaders/acc.comp | 29 + .../src/ggml-vulkan/vulkan-shaders/add.comp | 29 + .../ggml-vulkan/vulkan-shaders/argsort.comp | 69 + .../src/ggml-vulkan/vulkan-shaders/clamp.comp | 17 + .../ggml-vulkan/vulkan-shaders/concat.comp | 41 + .../vulkan-shaders/contig_copy.comp | 42 + .../src/ggml-vulkan/vulkan-shaders/copy.comp | 20 + .../src/ggml-vulkan/vulkan-shaders/cos.comp | 17 + .../vulkan-shaders/dequant_f32.comp | 20 + .../vulkan-shaders/dequant_funcs.comp | 118 + .../vulkan-shaders/dequant_funcs_cm2.comp | 325 + .../vulkan-shaders/dequant_head.comp | 13 + .../vulkan-shaders/dequant_iq4_nl.comp | 32 + .../vulkan-shaders/dequant_q2_k.comp | 34 + .../vulkan-shaders/dequant_q3_k.comp | 42 + .../vulkan-shaders/dequant_q4_0.comp | 30 + .../vulkan-shaders/dequant_q4_1.comp | 32 + .../vulkan-shaders/dequant_q4_k.comp | 68 + .../vulkan-shaders/dequant_q5_0.comp | 34 + .../vulkan-shaders/dequant_q5_1.comp | 35 + .../vulkan-shaders/dequant_q5_k.comp | 70 + .../vulkan-shaders/dequant_q6_k.comp | 33 + .../vulkan-shaders/dequant_q8_0.comp | 31 + .../vulkan-shaders/diag_mask_inf.comp | 34 + .../src/ggml-vulkan/vulkan-shaders/div.comp | 27 + .../vulkan-shaders/flash_attn_cm2.comp | 289 + .../src/ggml-vulkan/vulkan-shaders/gelu.comp | 25 + .../vulkan-shaders/gelu_quick.comp | 23 + .../vulkan-shaders/generic_binary_head.comp | 64 + .../vulkan-shaders/generic_head.comp | 9 + .../vulkan-shaders/generic_unary_head.comp | 56 + .../ggml-vulkan/vulkan-shaders/get_rows.comp | 28 + .../vulkan-shaders/get_rows_quant.comp | 39 + .../vulkan-shaders/group_norm.comp | 66 + .../ggml-vulkan/vulkan-shaders/im2col.comp | 87 + .../vulkan-shaders/leaky_relu.comp | 22 + .../src/ggml-vulkan/vulkan-shaders/mul.comp | 27 + .../mul_mat_split_k_reduce.comp | 48 + .../vulkan-shaders/mul_mat_vec.comp | 152 + .../vulkan-shaders/mul_mat_vec_base.comp | 118 + .../vulkan-shaders/mul_mat_vec_nc.comp | 71 + .../vulkan-shaders/mul_mat_vec_p021.comp | 73 + .../vulkan-shaders/mul_mat_vec_q2_k.comp | 115 + .../vulkan-shaders/mul_mat_vec_q3_k.comp | 103 + .../vulkan-shaders/mul_mat_vec_q4_k.comp | 133 + .../vulkan-shaders/mul_mat_vec_q5_k.comp | 162 + .../vulkan-shaders/mul_mat_vec_q6_k.comp | 112 + .../ggml-vulkan/vulkan-shaders/mul_mm.comp | 631 ++ .../vulkan-shaders/mul_mm_cm2.comp | 328 + .../src/ggml-vulkan/vulkan-shaders/norm.comp | 44 + .../src/ggml-vulkan/vulkan-shaders/pad.comp | 28 + .../ggml-vulkan/vulkan-shaders/pool2d.comp | 74 + .../src/ggml-vulkan/vulkan-shaders/relu.comp | 21 + .../ggml-vulkan/vulkan-shaders/repeat.comp | 26 + .../ggml-vulkan/vulkan-shaders/rms_norm.comp | 42 + .../ggml-vulkan/vulkan-shaders/rope_head.comp | 49 + .../ggml-vulkan/vulkan-shaders/rope_neox.comp | 37 + .../ggml-vulkan/vulkan-shaders/rope_norm.comp | 37 + .../src/ggml-vulkan/vulkan-shaders/scale.comp | 24 + .../src/ggml-vulkan/vulkan-shaders/silu.comp | 22 + .../src/ggml-vulkan/vulkan-shaders/sin.comp | 17 + .../ggml-vulkan/vulkan-shaders/soft_max.comp | 174 + .../ggml-vulkan/vulkan-shaders/square.comp | 17 + .../ggml-vulkan/vulkan-shaders/sum_rows.comp | 37 + .../src/ggml-vulkan/vulkan-shaders/tanh.comp | 20 + .../vulkan-shaders/test_coopmat2_support.comp | 7 + .../vulkan-shaders/timestep_embedding.comp | 41 + .../src/ggml-vulkan/vulkan-shaders/types.comp | 323 + .../ggml-vulkan/vulkan-shaders/upscale.comp | 36 + .../vulkan-shaders/vulkan-shaders-gen.cpp | 594 ++ .../src/ggml-vulkan/vulkan-shaders/wkv6.comp | 87 + 75 files changed, 14627 insertions(+), 1 deletion(-) create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp diff --git a/Makefile.sync b/Makefile.sync index 3001487de..78333fd62 100644 --- a/Makefile.sync +++ b/Makefile.sync @@ -32,7 +32,7 @@ PATCHES=$(wildcard llama/patches/*.patch) apply-patches: $(addsuffix ed, $(PATCHES)) %.patched: %.patch - @if git -c user.name=nobody -c 'user.email=<>' -C $(WORKDIR) am -3 $(realpath $<); then touch $@; else git -C $(WORKDIR) am --abort; exit 1; fi + @if git -c commit.gpgSign=false -c user.name=nobody -c 'user.email=<>' -C $(WORKDIR) am -3 $(realpath $<); then touch $@; else git -C $(WORKDIR) am --abort; exit 1; fi .PHONY: checkout checkout: $(WORKDIR) diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt b/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt new file mode 100644 index 000000000..9501de736 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt @@ -0,0 +1,92 @@ +find_package(Vulkan COMPONENTS glslc REQUIRED) + +if (Vulkan_FOUND) + message(STATUS "Vulkan found") + + ggml_add_backend_library(ggml-vulkan + ggml-vulkan.cpp + ../../include/ggml-vulkan.h + ) + + # Compile a test shader to determine whether GL_NV_cooperative_matrix2 is supported. + # If it's not, there will be an error to stderr. + # If it's supported, set a define to indicate that we should compile those shaders + execute_process(COMMAND ${Vulkan_GLSLC_EXECUTABLE} -o - -fshader-stage=compute --target-env=vulkan1.3 "${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders/test_coopmat2_support.comp" + OUTPUT_VARIABLE glslc_output + ERROR_VARIABLE glslc_error) + + if (${glslc_error} MATCHES ".*extension not supported: GL_NV_cooperative_matrix2.*") + message(STATUS "GL_NV_cooperative_matrix2 not supported by glslc") + else() + message(STATUS "GL_NV_cooperative_matrix2 supported by glslc") + add_compile_definitions(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + endif() + + target_link_libraries(ggml-vulkan PRIVATE Vulkan::Vulkan) + target_include_directories(ggml-vulkan PRIVATE ${CMAKE_CURRENT_BINARY_DIR}) + + # Workaround to the "can't dereference invalidated vector iterator" bug in clang-cl debug build + # Posssibly relevant: https://stackoverflow.com/questions/74748276/visual-studio-no-displays-the-correct-length-of-stdvector + if (MSVC AND CMAKE_CXX_COMPILER_ID STREQUAL "Clang") + add_compile_definitions(_ITERATOR_DEBUG_LEVEL=0) + endif() + + if (GGML_VULKAN_CHECK_RESULTS) + add_compile_definitions(GGML_VULKAN_CHECK_RESULTS) + endif() + + if (GGML_VULKAN_DEBUG) + add_compile_definitions(GGML_VULKAN_DEBUG) + endif() + + if (GGML_VULKAN_MEMORY_DEBUG) + add_compile_definitions(GGML_VULKAN_MEMORY_DEBUG) + endif() + + if (GGML_VULKAN_SHADER_DEBUG_INFO) + add_compile_definitions(GGML_VULKAN_SHADER_DEBUG_INFO) + endif() + + if (GGML_VULKAN_PERF) + add_compile_definitions(GGML_VULKAN_PERF) + endif() + + if (GGML_VULKAN_VALIDATE) + add_compile_definitions(GGML_VULKAN_VALIDATE) + endif() + + if (GGML_VULKAN_RUN_TESTS) + add_compile_definitions(GGML_VULKAN_RUN_TESTS) + endif() + + add_subdirectory(vulkan-shaders) + + set (_ggml_vk_genshaders_cmd vulkan-shaders-gen) + set (_ggml_vk_header ${CMAKE_CURRENT_BINARY_DIR}/ggml-vulkan-shaders.hpp) + set (_ggml_vk_source ${CMAKE_CURRENT_BINARY_DIR}/ggml-vulkan-shaders.cpp) + set (_ggml_vk_input_dir ${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders) + set (_ggml_vk_output_dir ${CMAKE_CURRENT_BINARY_DIR}/vulkan-shaders.spv) + + file(GLOB _ggml_vk_shader_deps "${_ggml_vk_input_dir}/*.comp") + + add_custom_command( + OUTPUT ${_ggml_vk_header} + ${_ggml_vk_source} + + COMMAND "$/${_ggml_vk_genshaders_cmd}" + --glslc ${Vulkan_GLSLC_EXECUTABLE} + --input-dir ${_ggml_vk_input_dir} + --output-dir ${_ggml_vk_output_dir} + --target-hpp ${_ggml_vk_header} + --target-cpp ${_ggml_vk_source} + --no-clean + + DEPENDS ${_ggml_vk_shader_deps} ${_ggml_vk_genshaders_cmd} + COMMENT "Generate vulkan shaders" + ) + + target_sources(ggml-vulkan PRIVATE ${_ggml_vk_source} ${_ggml_vk_header}) + +else() + message(WARNING "Vulkan not found") +endif() diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp new file mode 100644 index 000000000..d75cd6d61 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -0,0 +1,8745 @@ +#include "ggml-vulkan.h" +#include +#if defined(GGML_VULKAN_RUN_TESTS) || defined(GGML_VULKAN_PERF) || defined(GGML_VULKAN_CHECK_RESULTS) +#include +#include "ggml-cpu.h" +#endif + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ggml-impl.h" +#include "ggml-backend-impl.h" + +#include "ggml-vulkan-shaders.hpp" + +#define VK_API_VERSION VK_API_VERSION_1_2 + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +#define VK_VENDOR_ID_AMD 0x1002 +#define VK_VENDOR_ID_APPLE 0x106b +#define VK_VENDOR_ID_INTEL 0x8086 +#define VK_VENDOR_ID_NVIDIA 0x10de + +#define VK_DEVICE_DESCRIPTOR_POOL_SIZE 32 + +#define GGML_VK_MAX_NODES 8192 + +#define MAX_VK_BUFFERS 256 + +#define VK_CHECK(err, msg) \ + do { \ + vk::Result err_ = (err); \ + if (err_ != vk::Result::eSuccess) { \ + fprintf(stderr, "ggml_vulkan: %s error %s at %s:%d\n", \ + #err, to_string(err_).c_str(), __FILE__, __LINE__); \ + exit(1); \ + } \ + } while (0) + +#ifdef GGML_VULKAN_DEBUG +#define VK_LOG_DEBUG(msg) std::cerr << msg << std::endl +#else +#define VK_LOG_DEBUG(msg) ((void) 0) +#endif // GGML_VULKAN_DEBUG + +struct ggml_backend_vk_context; + +struct vk_queue { + uint32_t queue_family_index; + vk::Queue queue; + vk::CommandPool pool; + uint32_t cmd_buffer_idx; + std::vector cmd_buffers; + + vk::PipelineStageFlags stage_flags; + + bool transfer_only; +}; + +struct vk_pipeline_struct { + std::string name; + vk::ShaderModule shader_module; + vk::DescriptorSetLayout dsl; + std::vector descriptor_pools; + std::vector descriptor_sets; + uint32_t descriptor_set_idx; + vk::PipelineLayout layout; + vk::Pipeline pipeline; + uint32_t push_constant_size; + uint32_t parameter_count; + std::array wg_denoms; + uint32_t align; +}; + +typedef std::shared_ptr vk_pipeline; +typedef std::weak_ptr vk_pipeline_ref; + +static void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline); + +struct vk_matmul_pipeline_struct { + vk_pipeline l, m, s; + vk_pipeline a_l, a_m, a_s; +}; + +typedef std::shared_ptr vk_matmul_pipeline; + +struct vk_matmul_pipeline2 { + vk_matmul_pipeline2() { + f16acc = std::make_shared(); + f32acc = std::make_shared(); + } + vk_matmul_pipeline f32acc; + vk_matmul_pipeline f16acc; +}; + +struct vk_device_struct; +typedef std::shared_ptr vk_device; +typedef std::weak_ptr vk_device_ref; + +struct vk_buffer_struct; +typedef std::shared_ptr vk_buffer; +typedef std::weak_ptr vk_buffer_ref; + +struct ggml_backend_vk_buffer_type_context { + std::string name; + vk_device device; +}; + +static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft); +static ggml_backend_buffer_t ggml_backend_vk_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size); +static size_t ggml_backend_vk_buffer_type_get_alignment(ggml_backend_buffer_type_t buft); +static size_t ggml_backend_vk_buffer_type_get_max_size(ggml_backend_buffer_type_t buft); +static size_t ggml_backend_vk_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor); +static ggml_backend_buffer_type_i ggml_backend_vk_buffer_type_interface = { + /* .get_name = */ ggml_backend_vk_buffer_type_name, + /* .alloc_buffer = */ ggml_backend_vk_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_vk_buffer_type_get_alignment, + /* .get_max_size = */ ggml_backend_vk_buffer_type_get_max_size, + /* .get_alloc_size = */ ggml_backend_vk_buffer_type_get_alloc_size, + /* .is_host = */ NULL, +}; + +#ifdef GGML_VULKAN_MEMORY_DEBUG +class vk_memory_logger; +#endif +#ifdef GGML_VULKAN_PERF +class vk_perf_logger; +#endif +static void ggml_vk_destroy_buffer(vk_buffer& buf); + +static constexpr uint32_t mul_mat_vec_max_cols = 8; + +struct vk_device_struct { + std::mutex mutex; + + vk::PhysicalDevice physical_device; + vk::PhysicalDeviceProperties properties; + std::string name; + uint64_t max_memory_allocation_size; + bool fp16; + bool pipeline_robustness; + vk::Device device; + uint32_t vendor_id; + vk_queue compute_queue; + vk_queue transfer_queue; + bool single_queue; + uint32_t subgroup_size; + uint32_t shader_core_count; + bool uma; + bool float_controls_rte_fp16; + + bool subgroup_size_control; + uint32_t subgroup_min_size; + uint32_t subgroup_max_size; + bool subgroup_require_full_support; + + bool coopmat_support; + bool coopmat_acc_f32_support; + bool coopmat_acc_f16_support; + uint32_t coopmat_m; + uint32_t coopmat_n; + uint32_t coopmat_k; + bool coopmat2; + + size_t idx; + + bool mul_mat_l; + bool mul_mat_m; + bool mul_mat_s; + bool mul_mat_id_l; + bool mul_mat_id_m; + bool mul_mat_id_s; + + vk_matmul_pipeline pipeline_matmul_f32; + vk_matmul_pipeline pipeline_matmul_f32_f16; + vk_matmul_pipeline2 pipeline_matmul_f16; + vk_matmul_pipeline2 pipeline_matmul_f16_f32; + vk_pipeline pipeline_matmul_split_k_reduce; + + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_COUNT]; + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat[GGML_TYPE_COUNT]; + + vk_matmul_pipeline pipeline_matmul_id_f32; + vk_matmul_pipeline2 pipeline_matmul_id_f16; + vk_matmul_pipeline2 pipeline_matmul_id_f16_f32; + + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id[GGML_TYPE_COUNT]; + + vk_pipeline pipeline_dequant[GGML_TYPE_COUNT]; + vk_pipeline pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_COUNT][mul_mat_vec_max_cols]; + vk_pipeline pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_COUNT][mul_mat_vec_max_cols]; + vk_pipeline pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_COUNT]; + + vk_pipeline pipeline_mul_mat_vec_p021_f16_f32; + vk_pipeline pipeline_mul_mat_vec_nc_f16_f32; + vk_pipeline pipeline_get_rows[GGML_TYPE_COUNT]; + vk_pipeline pipeline_get_rows_f32[GGML_TYPE_COUNT]; + vk_pipeline pipeline_acc_f32; + vk_pipeline pipeline_add_f32, pipeline_add_f32_norepeat; + vk_pipeline pipeline_add_f16_f32_f16, pipeline_add_f16_f32_f16_norepeat; + vk_pipeline pipeline_mul_f32, pipeline_mul_f32_norepeat; + vk_pipeline pipeline_div_f32, pipeline_div_f32_norepeat; + vk_pipeline pipeline_concat_f32, pipeline_concat_f16, pipeline_concat_i32; + vk_pipeline pipeline_upscale_f32; + vk_pipeline pipeline_scale_f32; + vk_pipeline pipeline_sqr_f32; + vk_pipeline pipeline_sin_f32; + vk_pipeline pipeline_cos_f32; + vk_pipeline pipeline_clamp_f32; + vk_pipeline pipeline_pad_f32; + vk_pipeline pipeline_repeat_f32; + vk_pipeline pipeline_cpy_f32_f32, pipeline_cpy_f32_f16, pipeline_cpy_f16_f16; + vk_pipeline pipeline_contig_cpy_f32_f32, pipeline_contig_cpy_f32_f16, pipeline_contig_cpy_f16_f16; + vk_pipeline pipeline_norm_f32; + vk_pipeline pipeline_group_norm_f32; + vk_pipeline pipeline_rms_norm_f32; + vk_pipeline pipeline_gelu_f32; + vk_pipeline pipeline_gelu_quick_f32; + vk_pipeline pipeline_silu_f32; + vk_pipeline pipeline_relu_f32; + vk_pipeline pipeline_leaky_relu_f32; + vk_pipeline pipeline_tanh_f32; + vk_pipeline pipeline_diag_mask_inf_f32; + vk_pipeline pipeline_soft_max_f32, pipeline_soft_max_f32_f16; + vk_pipeline pipeline_soft_max_f32_wg512, pipeline_soft_max_f32_f16_wg512; + vk_pipeline pipeline_rope_norm_f32, pipeline_rope_norm_f16; + vk_pipeline pipeline_rope_neox_f32, pipeline_rope_neox_f16; + vk_pipeline pipeline_argsort_f32; + vk_pipeline pipeline_sum_rows_f32; + vk_pipeline pipeline_im2col_f32, pipeline_im2col_f32_f16; + vk_pipeline pipeline_timestep_embedding_f32; + vk_pipeline pipeline_pool2d_f32; + vk_pipeline pipeline_rwkv_wkv6_f32; + + // [2][2][2] is for {f16acc,f32acc}x{large,small_rows}x{unaligned, aligned} + vk_pipeline pipeline_flash_attn_f32_f16_D64[GGML_TYPE_COUNT][2][2][2]; + vk_pipeline pipeline_flash_attn_f32_f16_D80[GGML_TYPE_COUNT][2][2][2]; + vk_pipeline pipeline_flash_attn_f32_f16_D96[GGML_TYPE_COUNT][2][2][2]; + vk_pipeline pipeline_flash_attn_f32_f16_D112[GGML_TYPE_COUNT][2][2][2]; + vk_pipeline pipeline_flash_attn_f32_f16_D128[GGML_TYPE_COUNT][2][2][2]; + vk_pipeline pipeline_flash_attn_f32_f16_D256[GGML_TYPE_COUNT][2][2][2]; + + std::unordered_map pipelines; + std::unordered_map pipeline_descriptor_set_requirements; + + std::vector> pinned_memory; + + vk::Fence fence; + vk_buffer sync_staging; + + ggml_backend_buffer_type buffer_type; + +#ifdef GGML_VULKAN_MEMORY_DEBUG + std::unique_ptr memory_logger; +#endif +#ifdef GGML_VULKAN_PERF + std::unique_ptr perf_logger; +#endif + + ~vk_device_struct() { + VK_LOG_DEBUG("destroy device " << name); + + device.destroyFence(fence); + + ggml_vk_destroy_buffer(sync_staging); + + device.destroyCommandPool(compute_queue.pool); + if (!single_queue) { + device.destroyCommandPool(transfer_queue.pool); + } + + for (auto& pipeline : pipelines) { + if (pipeline.second.expired()) { + continue; + } + + vk_pipeline pl = pipeline.second.lock(); + ggml_vk_destroy_pipeline(device, pl); + } + pipelines.clear(); + + device.destroy(); + } +}; + +struct vk_buffer_struct { + vk::Buffer buffer = VK_NULL_HANDLE; + vk::DeviceMemory device_memory = VK_NULL_HANDLE; + vk::MemoryPropertyFlags memory_property_flags; + void * ptr; + size_t size = 0; + + vk_device device; + + ~vk_buffer_struct() { + if (size == 0) { + return; + } + VK_LOG_DEBUG("~vk_buffer_struct(" << buffer << ", " << size << ")"); + + device->device.freeMemory(device_memory); + device->device.destroyBuffer(buffer); + } +}; + +struct vk_subbuffer { + vk_buffer buffer; + uint64_t offset; + uint64_t size; + + operator vk::DescriptorBufferInfo() const { + return { buffer->buffer, offset, size }; + } +}; + +struct vk_semaphore { + vk::Semaphore s; + uint64_t value; +}; + +struct vk_submission { + vk::CommandBuffer buffer; + std::vector wait_semaphores; + std::vector signal_semaphores; +}; + +typedef std::vector vk_sequence; + +struct vk_mat_mat_push_constants { + uint32_t M; uint32_t N; uint32_t K; + uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; + uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; + uint32_t k_split; + uint32_t ne02; uint32_t ne12; uint32_t broadcast2; uint32_t broadcast3; +}; +struct vk_mat_vec_push_constants { + uint32_t ncols; uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; + uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; + uint32_t ne02; uint32_t ne12; uint32_t broadcast2; uint32_t broadcast3; +}; + +struct vk_mat_mat_id_push_constants { + uint32_t M; uint32_t N; uint32_t K; + uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; + uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; + uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11; +}; +struct vk_mat_vec_id_push_constants { + uint32_t ncols; uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; + uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; + uint32_t nei0; uint32_t ne11; +}; + +struct vk_flash_attn_push_constants { + uint32_t N; + uint32_t KV; + + uint32_t ne1; + uint32_t ne2; + uint32_t ne3; + + uint32_t neq2; + uint32_t neq3; + uint32_t nek2; + uint32_t nek3; + uint32_t nev2; + uint32_t nev3; + uint32_t nem1; + + uint32_t nb02; + uint32_t nb03; + uint32_t nb12; + uint32_t nb13; + uint32_t nb22; + uint32_t nb23; + uint32_t nb31; + + float scale; + float max_bias; + float logit_softcap; + + uint32_t mask; + uint32_t n_head_log2; + float m0; + float m1; +}; + +struct vk_op_push_constants { + uint32_t KX; + uint32_t KY; + float param1; + float param2; +}; + +struct vk_op_unary_push_constants { + uint32_t ne; + uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; + uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; uint32_t nb10; uint32_t nb11; uint32_t nb12; uint32_t nb13; + uint32_t misalign_offsets; + float param1; float param2; + uint32_t ne0_012mp; uint32_t ne0_012L; + uint32_t ne0_01mp; uint32_t ne0_01L; + uint32_t ne0_0mp; uint32_t ne0_0L; + uint32_t ne1_012mp; uint32_t ne1_012L; + uint32_t ne1_01mp; uint32_t ne1_01L; + uint32_t ne1_0mp; uint32_t ne1_0L; +}; +static_assert(sizeof(vk_op_unary_push_constants) <= 128, "sizeof(vk_op_unary_push_constants) must be <= 128"); + +// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1. +// Precompute mp (m' in the paper) and L such that division +// can be computed using a multiply (high 32b of 64b result) +// and a shift: +// +// n/d = (mulhi(n, mp) + n) >> L; +static void init_fastdiv_values(uint32_t d, uint32_t &mp, uint32_t &L) +{ + // compute L = ceil(log2(d)); + L = 0; + while (L < 32 && (uint32_t{1} << L) < d) { + L++; + } + + mp = (uint32_t)((uint64_t{1} << 32) * ((uint64_t{1} << L) - d) / d + 1); +} + +template void init_pushconst_fastdiv(T &p) { + GGML_UNUSED(p); + static_assert(!std::is_const::value, "unexpected type"); +} + +template <> void init_pushconst_fastdiv(vk_op_unary_push_constants &p) { + // Compute magic values to divide by these six numbers. + init_fastdiv_values(p.ne02*p.ne01*p.ne00, p.ne0_012mp, p.ne0_012L); + init_fastdiv_values(p.ne01*p.ne00, p.ne0_01mp, p.ne0_01L); + init_fastdiv_values(p.ne00, p.ne0_0mp, p.ne0_0L); + init_fastdiv_values(p.ne12*p.ne11*p.ne10, p.ne1_012mp, p.ne1_012L); + init_fastdiv_values(p.ne11*p.ne10, p.ne1_01mp, p.ne1_01L); + init_fastdiv_values(p.ne10, p.ne1_0mp, p.ne1_0L); +} + +struct vk_op_binary_push_constants { + uint32_t ne; + uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; + uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; uint32_t nb10; uint32_t nb11; uint32_t nb12; uint32_t nb13; + uint32_t ne20; uint32_t ne21; uint32_t ne22; uint32_t ne23; uint32_t nb20; uint32_t nb21; uint32_t nb22; uint32_t nb23; + uint32_t misalign_offsets; + float param1; float param2; int32_t param3; +}; + +struct vk_op_diag_mask_push_constants { + uint32_t ncols; + uint32_t rows_per_channel; + int32_t n_past; +}; + +struct vk_op_rope_push_constants { + uint32_t ncols; + uint32_t n_dims; + float freq_scale; + uint32_t p_delta_rows; + float freq_base; + float ext_factor; + float attn_factor; + float corr_dims[2]; + float theta_scale; + uint32_t has_ff; +}; + +struct vk_op_soft_max_push_constants { + uint32_t KX; + uint32_t KY; + float scale; + float max_bias; + float m0; + float m1; + uint32_t n_head_log2; + uint32_t nrows_x; +}; + +struct vk_op_argsort_push_constants { + uint32_t ncols; + uint32_t ncols_pad; + int32_t order; +}; + +struct vk_op_im2col_push_constants { + uint32_t batch_offset; uint32_t offset_delta; + uint32_t IC; + uint32_t IW; uint32_t IH; + uint32_t OW; uint32_t OH; + uint32_t KW; uint32_t KH; + uint32_t pelements; + uint32_t CHW; + int32_t s0; int32_t s1; + int32_t p0; int32_t p1; + int32_t d0; int32_t d1; +}; + +struct vk_op_timestep_embedding_push_constants { + uint32_t nb1; + uint32_t dim; + uint32_t max_period; +}; + +struct vk_op_pool2d_push_constants { + uint32_t IW; uint32_t IH; + uint32_t OW; uint32_t OH; + uint32_t OC; + uint32_t pelements; + uint32_t op; + int32_t k0; int32_t k1; + int32_t s0; int32_t s1; + int32_t p0; int32_t p1; +}; + +struct vk_op_rwkv_wkv6_push_constants { + uint32_t B; + uint32_t T; + uint32_t C; + uint32_t H; +}; + +// Allow pre-recording command buffers +struct vk_staging_memcpy { + vk_staging_memcpy(void * _dst, const void * _src, size_t _n) : dst(_dst), src(_src), n(_n) {} + + void * dst; + const void * src; + size_t n; +}; + +struct vk_op_upscale_push_constants { + uint32_t ne; uint32_t a_offset; uint32_t d_offset; + uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; + uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; + float sf0; float sf1; float sf2; float sf3; +}; + +struct vk_context_struct { + vk_submission * s; + std::vector seqs; + + int exit_tensor_idx; + + std::vector in_memcpys; + std::vector out_memcpys; + + vk_queue * q; +}; +typedef std::shared_ptr vk_context; +typedef std::weak_ptr vk_context_ref; + +struct ggml_vk_garbage_collector { + std::vector tl_semaphores; + std::vector semaphores; + std::vector events; + std::vector temp_buffers; + std::vector contexts; +}; + +#if defined(GGML_VULKAN_MEMORY_DEBUG) || defined(GGML_VULKAN_DEBUG) +#define VK_LOG_MEMORY(msg) std::cerr << "ggml_vulkan memory: " << msg << std::endl + +static std::string format_size(size_t size) { + const size_t kib = 1024; + const size_t mib = kib * 1024; + const size_t gib = mib * 1024; + + std::ostringstream oss; + oss << std::fixed << std::setprecision(2); + + if (size >= gib) { + oss << static_cast(size) / gib << " GiB"; + } else if (size >= mib) { + oss << static_cast(size) / mib << " MiB"; + } else if (size >= kib) { + oss << static_cast(size) / kib << " KiB"; + } else { + oss << size << " B"; + } + + return oss.str(); +} + +static std::mutex log_mutex; + +class vk_memory_logger { +public: + vk_memory_logger(): total_device(0), total_host(0) {} + void log_allocation(vk_buffer_ref buf_ref, size_t size); + void log_deallocation(vk_buffer_ref buf_ref); + +private: + std::map allocations; // Track allocations + size_t total_device; + size_t total_host; +}; +#else +#define VK_LOG_MEMORY(msg) ((void) 0) +#endif // GGML_VULKAN_MEMORY_DEBUG + +#if defined(GGML_VULKAN_PERF) + +class vk_perf_logger { +public: + void print_timings() { + std::cerr << "----------------\nVulkan Timings:" << std::endl; + for (const auto& t : timings) { + uint64_t total = 0; + for (const auto& time : t.second) { + total += time; + } + std::cerr << t.first << ": " << t.second.size() << " x " << (total / t.second.size() / 1000.0) << " ms" << std::endl; + } + + timings.clear(); + } + + void log_timing(const ggml_tensor * node, uint64_t time) { + if (node->op == GGML_OP_UNARY) { + timings[ggml_unary_op_name(ggml_get_unary_op(node))].push_back(time); + return; + } + if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { + const uint64_t m = node->src[0]->ne[1]; + const uint64_t n = node->src[1]->ne[1]; + const uint64_t k = node->src[1]->ne[0]; + std::string name = ggml_op_name(node->op); + if (n == 1) { + name += "_VEC m=" + std::to_string(m) + " k=" + std::to_string(k); + } else { + name += " m=" + std::to_string(m) + " n=" + std::to_string(n) + " k=" + std::to_string(k); + } + timings[name].push_back(time); + return; + } + timings[ggml_op_name(node->op)].push_back(time); + } +private: + std::map> timings; +}; +#endif // GGML_VULKAN_PERF + +struct ggml_backend_vk_context { + std::string name; + + vk_device device; + + size_t semaphore_idx, event_idx; + ggml_vk_garbage_collector gc; + size_t prealloc_size_x, prealloc_size_y, prealloc_size_split_k; + vk_buffer prealloc_x, prealloc_y, prealloc_split_k; + vk::Fence fence; + + vk_buffer buffer_pool[MAX_VK_BUFFERS]; + + vk_context_ref compute_ctx; + vk_context_ref transfer_ctx; + + std::vector tensor_ctxs; +}; + +static void * const vk_ptr_base = (void *)(uintptr_t) 0x1000; // NOLINT + +static uint64_t vk_tensor_offset(const ggml_tensor * tensor) { + if (tensor->view_src) { + return (uint8_t *) tensor->view_src->data - (uint8_t *) vk_ptr_base; + } + return (uint8_t *) tensor->data - (uint8_t *) vk_ptr_base; +} + +struct ggml_backend_vk_buffer_context { + vk_device_ref device; + vk_buffer dev_buffer; + std::string name; + + ggml_backend_vk_buffer_context(vk_device_ref device, vk_buffer&& dev_buffer, std::string& name) : + device(device), + dev_buffer(dev_buffer), + name(name) { + } + + ~ggml_backend_vk_buffer_context() { + ggml_vk_destroy_buffer(dev_buffer); + } +}; + +#ifdef GGML_VULKAN_MEMORY_DEBUG +void vk_memory_logger::log_allocation(vk_buffer_ref buf_ref, size_t size) { + std::lock_guard guard(log_mutex); + vk_buffer buf = buf_ref.lock(); + const bool device = bool(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eDeviceLocal); + const std::string type = device ? "device" : "host"; + allocations[buf->buffer] = size; + total_device += device ? size : 0; + total_host += device ? 0 : size; + VK_LOG_MEMORY(buf->device->name << ": +" << format_size(size) << " " << type << " at " << buf->buffer << ". Total device: " << format_size(total_device) << ", total host: " << format_size(total_host)); +} + +void vk_memory_logger::log_deallocation(vk_buffer_ref buf_ref) { + if (buf_ref.expired() || buf_ref.lock()->size == 0) { + return; + } + + std::lock_guard guard(log_mutex); + vk_buffer buf = buf_ref.lock(); + const bool device = bool(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eDeviceLocal); + std::string type = device ? "device" : "host"; + auto it = allocations.find(buf->buffer); + total_device -= device ? it->second : 0; + total_host -= device ? 0 : it->second; + if (it != allocations.end()) { + VK_LOG_MEMORY(buf->device->name << ": -" << format_size(it->second) << " " << type << " at " << buf->buffer << ". Total device: " << format_size(total_device) << ", total host: " << format_size(total_host)); + allocations.erase(it); + } else { + VK_LOG_MEMORY("ERROR " << buf->device->name << ": Attempted to deallocate unknown " << type << " memory at " << buf->buffer); + } +} +#endif // GGML_VULKAN_MEMORY_DEBUG + +struct vk_instance_t { + vk::Instance instance; + + std::vector device_indices; + vk_device devices[GGML_VK_MAX_DEVICES]; +}; + +static bool vk_instance_initialized = false; +static vk_instance_t vk_instance; + +#ifdef GGML_VULKAN_CHECK_RESULTS +static size_t vk_skip_checks; +static size_t vk_output_tensor; + +static void ggml_vk_print_tensor(const ggml_tensor * tensor, const char * name); +static void ggml_vk_check_results_0(ggml_tensor * tensor); +static void ggml_vk_check_results_1(ggml_tensor * tensor); +#endif + +typedef void (*ggml_vk_func_t)(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); + +static void ggml_backend_vk_free(ggml_backend_t backend); + +// variables to track number of compiles in progress +static uint32_t compile_count = 0; +static std::mutex compile_count_mutex; +static std::condition_variable compile_count_cond; + +static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipeline, const std::string name, size_t spv_size, const void* spv_data, const std::string entrypoint, + uint32_t parameter_count, uint32_t push_constant_size, std::array wg_denoms, std::vector specialization_constants, + uint32_t align, bool disable_robustness, bool require_full_subgroups, uint32_t required_subgroup_size) { + VK_LOG_DEBUG("ggml_vk_create_pipeline(" << device->name << ", " << name << ", " << entrypoint << ", " << parameter_count << ", " << push_constant_size << + ", (" << wg_denoms[0] << "," << wg_denoms[1] << "," << wg_denoms[2] << "), specialization_constants, " << align << + ", " << disable_robustness << ", " << require_full_subgroups << ", " << required_subgroup_size << ")"); + GGML_ASSERT(parameter_count > 0); + GGML_ASSERT(wg_denoms[0] > 0 && wg_denoms[1] > 0 && wg_denoms[2] > 0); // NOLINT + + pipeline = std::make_shared(); + pipeline->name = name; + pipeline->parameter_count = parameter_count; + pipeline->push_constant_size = push_constant_size; + pipeline->wg_denoms = wg_denoms; + pipeline->align = align; + + vk::ShaderModuleCreateInfo shader_module_create_info({}, spv_size, reinterpret_cast(spv_data)); + pipeline->shader_module = device->device.createShaderModule(shader_module_create_info); + + std::vector dsl_binding; + std::vector dsl_binding_flags; + for (uint32_t i = 0; i < parameter_count; i++) { + dsl_binding.push_back({i, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute}); + dsl_binding_flags.push_back({}); + } + + vk::DescriptorSetLayoutBindingFlagsCreateInfo dslbfci = { dsl_binding_flags }; + + vk::PushConstantRange pcr( + vk::ShaderStageFlagBits::eCompute, + 0, + pipeline->push_constant_size + ); + + vk::DescriptorSetLayoutCreateInfo descriptor_set_layout_create_info( + {}, + dsl_binding); + descriptor_set_layout_create_info.setPNext(&dslbfci); + pipeline->dsl = device->device.createDescriptorSetLayout(descriptor_set_layout_create_info); + + vk::DescriptorPoolSize descriptor_pool_size(vk::DescriptorType::eStorageBuffer, pipeline->parameter_count * VK_DEVICE_DESCRIPTOR_POOL_SIZE); + vk::DescriptorPoolCreateInfo descriptor_pool_create_info({}, VK_DEVICE_DESCRIPTOR_POOL_SIZE, descriptor_pool_size); + pipeline->descriptor_pools.push_back(device->device.createDescriptorPool(descriptor_pool_create_info)); + + pipeline->descriptor_set_idx = 0; + + vk::PipelineLayoutCreateInfo pipeline_layout_create_info(vk::PipelineLayoutCreateFlags(), pipeline->dsl, pcr); + pipeline->layout = device->device.createPipelineLayout(pipeline_layout_create_info); + + std::vector specialization_entries(specialization_constants.size()); + + for (size_t i = 0; i < specialization_constants.size(); i++) { + specialization_entries[i].constantID = i; + specialization_entries[i].offset = i * sizeof(uint32_t); + specialization_entries[i].size = sizeof(uint32_t); + } + + vk::SpecializationInfo specialization_info( + specialization_entries.size(), + specialization_entries.data(), + specialization_constants.size() * sizeof(uint32_t), + specialization_constants.data() + ); + + vk::PipelineShaderStageCreateFlags pipeline_shader_stage_create_flags{}; + + if (device->subgroup_require_full_support && require_full_subgroups) { + pipeline_shader_stage_create_flags |= vk::PipelineShaderStageCreateFlagBits::eRequireFullSubgroupsEXT; + } + + vk::PipelineShaderStageCreateInfo pipeline_shader_create_info( + pipeline_shader_stage_create_flags, + vk::ShaderStageFlagBits::eCompute, + pipeline->shader_module, + entrypoint.c_str(), + &specialization_info); + + vk::PipelineShaderStageRequiredSubgroupSizeCreateInfoEXT pipeline_shader_stage_required_subgroup_size_create_info; + pipeline_shader_stage_required_subgroup_size_create_info.requiredSubgroupSize = required_subgroup_size; + if (device->subgroup_size_control && required_subgroup_size > 0) { + GGML_ASSERT(device->subgroup_min_size <= required_subgroup_size && required_subgroup_size <= device->subgroup_max_size); + pipeline_shader_create_info.setPNext(&pipeline_shader_stage_required_subgroup_size_create_info); + } + + vk::ComputePipelineCreateInfo compute_pipeline_create_info( + vk::PipelineCreateFlags{}, + pipeline_shader_create_info, + pipeline->layout); + + vk::PipelineRobustnessCreateInfoEXT rci; + + if (device->pipeline_robustness && disable_robustness) { + rci.storageBuffers = vk::PipelineRobustnessBufferBehaviorEXT::eDisabled; + rci.uniformBuffers = vk::PipelineRobustnessBufferBehaviorEXT::eDisabled; + compute_pipeline_create_info.setPNext(&rci); + } + + pipeline->pipeline = device->device.createComputePipeline(VK_NULL_HANDLE, compute_pipeline_create_info).value; + + { + std::lock_guard guard(device->mutex); + device->pipelines.insert({ pipeline->name, pipeline }); + } + + { + std::lock_guard guard(compile_count_mutex); + assert(compile_count > 0); + compile_count--; + + // "Progress bar" for shader compiles + static uint32_t total_compile_count = 0; + if ((total_compile_count++ % 10) == 0) { + std::cerr << "."; + } + } + compile_count_cond.notify_all(); +} + +static void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline) { + VK_LOG_DEBUG("ggml_pipeline_destroy_pipeline(" << pipeline->name << ")"); + for (auto& pool : pipeline->descriptor_pools) { + device.destroyDescriptorPool(pool); + } + pipeline->descriptor_pools.clear(); + pipeline->descriptor_sets.clear(); + pipeline->descriptor_set_idx = 0; + + device.destroyDescriptorSetLayout(pipeline->dsl); + + device.destroyPipelineLayout(pipeline->layout); + + device.destroyShaderModule(pipeline->shader_module); + + device.destroyPipeline(pipeline->pipeline); +} + +static void ggml_pipeline_request_descriptor_sets(vk_device& device, vk_pipeline& pipeline, uint32_t n) { + VK_LOG_DEBUG("ggml_pipeline_request_descriptor_sets(" << pipeline->name << ", " << n << ")"); + device->pipeline_descriptor_set_requirements[pipeline->name] += n; +} + +static void ggml_pipeline_allocate_descriptor_sets(vk_device& device) { + std::lock_guard guard(device->mutex); + + for (auto& pair : device->pipeline_descriptor_set_requirements) { + vk_pipeline pipeline = device->pipelines.at(pair.first).lock(); + const uint64_t n = pair.second; + + VK_LOG_DEBUG("ggml_pipeline_allocate_descriptor_sets(" << pipeline->name << ", " << n << ")"); + + if (pipeline->descriptor_sets.size() >= pipeline->descriptor_set_idx + n) { + // Enough descriptors are available + continue; + } + + uint32_t to_alloc = pipeline->descriptor_set_idx + n - pipeline->descriptor_sets.size(); + uint32_t pool_remaining = VK_DEVICE_DESCRIPTOR_POOL_SIZE - pipeline->descriptor_sets.size() % VK_DEVICE_DESCRIPTOR_POOL_SIZE; + uint32_t pool_idx = pipeline->descriptor_sets.size() / VK_DEVICE_DESCRIPTOR_POOL_SIZE; + + while (to_alloc > 0) { + const uint32_t alloc_count = std::min(pool_remaining, to_alloc); + to_alloc -= alloc_count; + pool_remaining = VK_DEVICE_DESCRIPTOR_POOL_SIZE; + + if (pool_idx >= pipeline->descriptor_pools.size()) { + vk::DescriptorPoolSize descriptor_pool_size(vk::DescriptorType::eStorageBuffer, pipeline->parameter_count * VK_DEVICE_DESCRIPTOR_POOL_SIZE); + vk::DescriptorPoolCreateInfo descriptor_pool_create_info({}, VK_DEVICE_DESCRIPTOR_POOL_SIZE, descriptor_pool_size); + pipeline->descriptor_pools.push_back(device->device.createDescriptorPool(descriptor_pool_create_info)); + } + + std::vector layouts(alloc_count); + for (uint32_t i = 0; i < alloc_count; i++) { + layouts[i] = pipeline->dsl; + } + vk::DescriptorSetAllocateInfo descriptor_set_alloc_info(pipeline->descriptor_pools[pool_idx], alloc_count, layouts.data()); + std::vector sets = device->device.allocateDescriptorSets(descriptor_set_alloc_info); + pipeline->descriptor_sets.insert(pipeline->descriptor_sets.end(), sets.begin(), sets.end()); + + pool_idx++; + } + } +} + +static void ggml_pipeline_cleanup(vk_pipeline& pipeline) { + VK_LOG_DEBUG("ggml_pipeline_cleanup(" << pipeline->name << ")"); + pipeline->descriptor_set_idx = 0; +} + +static vk::CommandBuffer ggml_vk_create_cmd_buffer(vk_device& device, vk_queue& q) { + VK_LOG_DEBUG("ggml_vk_create_cmd_buffer()"); + std::lock_guard guard(device->mutex); + + if (q.cmd_buffers.size() > q.cmd_buffer_idx) { + // Reuse command buffer + return q.cmd_buffers[q.cmd_buffer_idx++]; + } + + vk::CommandBufferAllocateInfo command_buffer_alloc_info( + q.pool, + vk::CommandBufferLevel::ePrimary, + 1); + const std::vector cmd_buffers = device->device.allocateCommandBuffers(command_buffer_alloc_info); + auto buf = cmd_buffers.front(); + + q.cmd_buffers.push_back(buf); + q.cmd_buffer_idx++; + + return buf; +} + +static vk_submission ggml_vk_create_submission(vk_device& device, vk_queue& q, std::vector wait_semaphores, std::vector signal_semaphores) { + VK_LOG_DEBUG("ggml_vk_create_submission()"); + vk_submission s; + s.buffer = ggml_vk_create_cmd_buffer(device, q); + s.wait_semaphores = std::move(wait_semaphores); + s.signal_semaphores = std::move(signal_semaphores); + return s; +} + +static void ggml_vk_submit(vk_context& ctx, vk::Fence fence) { + if (ctx->seqs.empty()) { + if (fence) { + ctx->q->queue.submit({}, fence); + } + return; + } + VK_LOG_DEBUG("ggml_vk_submit(" << ctx << ", " << fence << ")"); + + std::vector> tl_wait_vals; + std::vector> tl_signal_vals; + std::vector> tl_wait_semaphores; + std::vector> tl_signal_semaphores; + std::vector tl_submit_infos; + std::vector submit_infos; + int idx = -1; + std::vector> stage_flags; + + size_t reserve = 0; + + for (const auto& sequence : ctx->seqs) { + reserve += sequence.size(); + } + + // Pre-reserve vectors to prevent reallocation, which invalidates pointers + tl_wait_semaphores.reserve(reserve); + tl_wait_vals.reserve(reserve); + tl_signal_semaphores.reserve(reserve); + tl_signal_vals.reserve(reserve); + tl_submit_infos.reserve(reserve); + submit_infos.reserve(reserve); + stage_flags.reserve(reserve); + + for (const auto& sequence : ctx->seqs) { + for (const auto& submission : sequence) { + stage_flags.push_back({}); + idx++; + tl_wait_vals.push_back({}); + tl_wait_semaphores.push_back({}); + tl_signal_vals.push_back({}); + tl_signal_semaphores.push_back({}); + for (size_t i = 0; i < submission.wait_semaphores.size(); i++) { + stage_flags[idx].push_back(ctx->q->stage_flags); + tl_wait_vals[idx].push_back(submission.wait_semaphores[i].value); + tl_wait_semaphores[idx].push_back(submission.wait_semaphores[i].s); + } + for (size_t i = 0; i < submission.signal_semaphores.size(); i++) { + tl_signal_vals[idx].push_back(submission.signal_semaphores[i].value); + tl_signal_semaphores[idx].push_back(submission.signal_semaphores[i].s); + } + tl_submit_infos.push_back({ + (uint32_t) submission.wait_semaphores.size(), + tl_wait_vals[idx].data(), + (uint32_t) submission.signal_semaphores.size(), + tl_signal_vals[idx].data(), + }); + tl_submit_infos[idx].sType = vk::StructureType::eTimelineSemaphoreSubmitInfo; + tl_submit_infos[idx].pNext = nullptr; + vk::SubmitInfo si{ + (uint32_t) submission.wait_semaphores.size(), + tl_wait_semaphores[idx].data(), + stage_flags[idx].data(), + 1, + &submission.buffer, + (uint32_t) submission.signal_semaphores.size(), + tl_signal_semaphores[idx].data(), + }; + si.setPNext(&tl_submit_infos[idx]); + submit_infos.push_back(si); + } + } + + ctx->q->queue.submit(submit_infos, fence); + + ctx->seqs.clear(); +} + +static uint32_t ggml_vk_find_queue_family_index(std::vector& queue_family_props, const vk::QueueFlags& required, const vk::QueueFlags& avoid, int32_t compute_index, uint32_t min_num_queues) { + VK_LOG_DEBUG("ggml_vk_find_queue_family_index()"); + const uint32_t qfsize = queue_family_props.size(); + + // Try with avoid preferences first + for (uint32_t i = 0; i < qfsize; i++) { + if (queue_family_props[i].queueCount >= min_num_queues && (compute_index < 0 || i != (uint32_t) compute_index) && queue_family_props[i].queueFlags & required && !(queue_family_props[i].queueFlags & avoid)) { + return i; + } + } + + // Fall back to only required + for (size_t i = 0; i < qfsize; i++) { + if (queue_family_props[i].queueCount >= min_num_queues && (compute_index < 0 || i != (uint32_t) compute_index) && queue_family_props[i].queueFlags & required) { + return i; + } + } + + // Fall back to reusing compute queue + for (size_t i = 0; i < qfsize; i++) { + if (queue_family_props[i].queueCount >= min_num_queues && queue_family_props[i].queueFlags & required) { + return i; + } + } + + // Fall back to ignoring min_num_queries + for (size_t i = 0; i < qfsize; i++) { + if (queue_family_props[i].queueFlags & required) { + return i; + } + } + + // All commands that are allowed on a queue that supports transfer operations are also allowed on a queue that supports either graphics or compute operations. + // Thus, if the capabilities of a queue family include VK_QUEUE_GRAPHICS_BIT or VK_QUEUE_COMPUTE_BIT, then reporting the VK_QUEUE_TRANSFER_BIT capability separately for that queue family is optional. + if (compute_index >= 0) { + return compute_index; + } + + std::cerr << "ggml_vulkan: No suitable queue family index found." << std::endl; + + for(auto &q_family : queue_family_props) { + std::cerr << "Queue number: " + std::to_string(q_family.queueCount) << " flags: " + to_string(q_family.queueFlags) << std::endl; + } + abort(); +} + +static void ggml_vk_create_queue(vk_device& device, vk_queue& q, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) { + VK_LOG_DEBUG("ggml_vk_create_queue()"); + std::lock_guard guard(device->mutex); + + q.queue_family_index = queue_family_index; + q.transfer_only = transfer_only; + + vk::CommandPoolCreateInfo command_pool_create_info_compute(vk::CommandPoolCreateFlags(VK_COMMAND_POOL_CREATE_TRANSIENT_BIT), queue_family_index); + q.pool = device->device.createCommandPool(command_pool_create_info_compute); + + q.cmd_buffer_idx = 0; + + q.queue = device->device.getQueue(queue_family_index, queue_index); + + q.stage_flags = stage_flags; +} + +static vk_context ggml_vk_create_context(ggml_backend_vk_context * ctx, vk_queue& q) { + vk_context result = std::make_shared(); + VK_LOG_DEBUG("ggml_vk_create_context(" << result << ")"); + ctx->gc.contexts.emplace_back(result); + result->q = &q; + return result; +} + +static vk_context ggml_vk_create_temporary_context(vk_queue& q) { + vk_context result = std::make_shared(); + VK_LOG_DEBUG("ggml_vk_create_temporary_context(" << result << ")"); + result->q = &q; + return result; +} + +static vk_semaphore * ggml_vk_create_binary_semaphore(ggml_backend_vk_context * ctx) { + VK_LOG_DEBUG("ggml_vk_create_timeline_semaphore()"); + vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eBinary, 0 }; + vk::SemaphoreCreateInfo ci{}; + ci.setPNext(&tci); + vk::Semaphore semaphore = ctx->device->device.createSemaphore(ci); + ctx->gc.semaphores.push_back({ semaphore, 0 }); + return &ctx->gc.semaphores[ctx->gc.semaphores.size() - 1]; +} + +static vk_semaphore * ggml_vk_create_timeline_semaphore(ggml_backend_vk_context * ctx) { + VK_LOG_DEBUG("ggml_vk_create_timeline_semaphore()"); + if (ctx->semaphore_idx >= ctx->gc.tl_semaphores.size()) { + vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eTimeline, 0 }; + vk::SemaphoreCreateInfo ci{}; + ci.setPNext(&tci); + vk::Semaphore semaphore = ctx->device->device.createSemaphore(ci); + ctx->gc.tl_semaphores.push_back({ semaphore, 0 }); + } + return &ctx->gc.tl_semaphores[ctx->semaphore_idx++]; +} + +static vk::Event ggml_vk_create_event(ggml_backend_vk_context * ctx) { + if (ctx->event_idx >= ctx->gc.events.size()) { + ctx->gc.events.push_back(ctx->device->device.createEvent({})); + } + return ctx->gc.events[ctx->event_idx++]; +} + +static void ggml_vk_queue_cleanup(vk_device& device, vk_queue& q) { + VK_LOG_DEBUG("ggml_vk_queue_cleanup()"); + std::lock_guard guard(device->mutex); + + // Requires command buffers to be done + device->device.resetCommandPool(q.pool); + q.cmd_buffer_idx = 0; +} + +static uint32_t find_properties(const vk::PhysicalDeviceMemoryProperties* mem_props, vk::MemoryRequirements* mem_req, vk::MemoryPropertyFlags flags) { + for (uint32_t i = 0; i < mem_props->memoryTypeCount; ++i) { + vk::MemoryType memory_type = mem_props->memoryTypes[i]; + if ((mem_req->memoryTypeBits & ((uint64_t)1 << i)) && + (flags & memory_type.propertyFlags) == flags && + mem_props->memoryHeaps[memory_type.heapIndex].size >= mem_req->size) { + return static_cast(i); + } + } + return UINT32_MAX; +} + +static vk_buffer ggml_vk_create_buffer(vk_device& device, size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags fallback_flags = vk::MemoryPropertyFlags(0)) { + VK_LOG_DEBUG("ggml_vk_create_buffer(" << device->name << ", " << size << ", " << to_string(req_flags) << ", " << to_string(fallback_flags) << ")"); + if (size > device->max_memory_allocation_size) { + throw vk::OutOfDeviceMemoryError("Requested buffer size exceeds device memory allocation limit"); + } + + std::lock_guard guard(device->mutex); + + vk_buffer buf = std::make_shared(); + + if (size == 0) { + buf->size = 0; + return buf; + } + + vk::BufferCreateInfo buffer_create_info{ + vk::BufferCreateFlags(), + size, + vk::BufferUsageFlagBits::eStorageBuffer | vk::BufferUsageFlagBits::eTransferSrc | vk::BufferUsageFlagBits::eTransferDst, + vk::SharingMode::eExclusive, + 0, + nullptr, + }; + + buf->buffer = device->device.createBuffer(buffer_create_info); + + vk::MemoryRequirements mem_req = device->device.getBufferMemoryRequirements(buf->buffer); + + vk::PhysicalDeviceMemoryProperties mem_props = device->physical_device.getMemoryProperties(); + + uint32_t memory_type_index = UINT32_MAX; + + memory_type_index = find_properties(&mem_props, &mem_req, req_flags); + buf->memory_property_flags = req_flags; + + if (memory_type_index == UINT32_MAX && fallback_flags) { + memory_type_index = find_properties(&mem_props, &mem_req, fallback_flags); + buf->memory_property_flags = fallback_flags; + } + + if (memory_type_index == UINT32_MAX) { + device->device.destroyBuffer(buf->buffer); + throw vk::OutOfDeviceMemoryError("No suitable memory type found"); + } + + try { + buf->device_memory = device->device.allocateMemory({ mem_req.size, memory_type_index }); + } catch (const vk::SystemError& e) { + if (buf->memory_property_flags != fallback_flags) { + // Try again with fallback flags + memory_type_index = find_properties(&mem_props, &mem_req, fallback_flags); + buf->memory_property_flags = fallback_flags; + + try { + buf->device_memory = device->device.allocateMemory({ mem_req.size, memory_type_index }); + } + catch (const vk::SystemError& e) { + device->device.destroyBuffer(buf->buffer); + throw e; + } + } else { + // Out of Host/Device memory, clean up buffer + device->device.destroyBuffer(buf->buffer); + throw e; + } + } + buf->ptr = nullptr; + + if (buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { + buf->ptr = device->device.mapMemory(buf->device_memory, 0, VK_WHOLE_SIZE); + } + + device->device.bindBufferMemory(buf->buffer, buf->device_memory, 0); + + buf->device = device; + buf->size = size; + +#ifdef GGML_VULKAN_MEMORY_DEBUG + device->memory_logger->log_allocation(buf, size); +#endif + + return buf; +} + +static vk_buffer ggml_vk_create_buffer_check(vk_device& device, size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags fallback_flags = vk::MemoryPropertyFlags(0)) { + try { + return ggml_vk_create_buffer(device, size, req_flags, fallback_flags); + } catch (const vk::SystemError& e) { + std::cerr << "ggml_vulkan: Memory allocation of size " << size << " failed." << std::endl; + std::cerr << "ggml_vulkan: " << e.what() << std::endl; + throw e; + } +} + +static vk_buffer ggml_vk_create_buffer_device(vk_device& device, size_t size) { + vk_buffer buf; + try { + if (device->uma) { + // Fall back to host memory type + buf = ggml_vk_create_buffer(device, size, vk::MemoryPropertyFlagBits::eDeviceLocal, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); + } else { + // use rebar if available, otherwise fallback to device only visible memory + buf = ggml_vk_create_buffer(device, size, vk::MemoryPropertyFlagBits::eDeviceLocal | vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent, vk::MemoryPropertyFlagBits::eDeviceLocal); + } + } catch (const vk::SystemError& e) { + std::cerr << "ggml_vulkan: Device memory allocation of size " << size << " failed." << std::endl; + std::cerr << "ggml_vulkan: " << e.what() << std::endl; + throw e; + } + + return buf; +} + +static void ggml_vk_destroy_buffer(vk_buffer& buf) { + if (buf == nullptr) { + return; + } + +#ifdef GGML_VULKAN_MEMORY_DEBUG + if (buf->device != nullptr) { + buf->device->memory_logger->log_deallocation(buf); + } +#endif + + buf.reset(); +} + +static vk_subbuffer ggml_vk_subbuffer(vk_buffer& buf) { + return { buf, 0, VK_WHOLE_SIZE }; +} + +static void ggml_vk_sync_buffers(vk_context& ctx) { + VK_LOG_DEBUG("ggml_vk_sync_buffers()"); + + const bool transfer_queue = ctx->q->transfer_only; + + ctx->s->buffer.pipelineBarrier( + ctx->q->stage_flags, + ctx->q->stage_flags, + {}, + { { + { !transfer_queue ? (vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) : (vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) }, + { !transfer_queue ? (vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) : (vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) } + } }, + {}, + {} + ); +} + +static void ggml_vk_wait_events(vk_context& ctx, std::vector&& events) { + VK_LOG_DEBUG("ggml_vk_wait_events()"); + if (events.empty()) { + return; + } + + ctx->s->buffer.waitEvents( + events, + ctx->q->stage_flags, + ctx->q->stage_flags, + {}, + {}, + {} + ); +} + +// number of rows/cols for flash attention shader +static constexpr uint32_t flash_attention_num_small_rows = 32; +static std::array fa_rows_cols(uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) { + GGML_UNUSED(clamp); + + // small rows, large cols + if (small_rows) { + return {flash_attention_num_small_rows, 128}; + } + // small cols to reduce register count + if (ggml_is_quantized(type) || D == 256) { + return {64, 32}; + } + return {64, 64}; +}; + +static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vector& warptile, bool mul_mat_id) { + // Needs to be kept up to date on shader changes + const uint32_t bank_conflict_offset = device->coopmat_support ? 8 : 1; + const uint32_t type_size = device->fp16 ? sizeof(ggml_fp16_t) : sizeof(float); + const uint32_t warps = warptile[0] / warptile[10]; + + const uint32_t load_bufs = (warptile[1] + warptile[2]) * (warptile[3] + bank_conflict_offset) * type_size; + const uint32_t mmid_row_ids = mul_mat_id ? 3072 * sizeof(uint32_t) : 0; + const uint32_t coopmat_stage = device->coopmat_support ? warptile[7] * warptile[8] / warps * sizeof(float) : 0; + + return (load_bufs + mmid_row_ids + coopmat_stage) <= device->properties.limits.maxComputeSharedMemorySize; +} + +static void ggml_vk_load_shaders(vk_device& device) { + VK_LOG_DEBUG("ggml_vk_load_shaders(" << device->name << ")"); + + std::cerr << "ggml_vulkan: Compiling shaders"; + + // some shaders have a minimum subgroup size + const uint32_t subgroup_size_16 = std::max(device->subgroup_size, 16u); + const uint32_t subgroup_size_32 = std::max(device->subgroup_size, 32u); + + // mulmat + std::vector l_warptile, m_warptile, s_warptile, + l_warptile_mmq, m_warptile_mmq, s_warptile_mmq, + l_warptile_mmq_k, m_warptile_mmq_k, s_warptile_mmq_k, + l_warptile_mmqid, m_warptile_mmqid, s_warptile_mmqid; + std::array l_wg_denoms, m_wg_denoms, s_wg_denoms, + l_mmq_wg_denoms, m_mmq_wg_denoms, s_mmq_wg_denoms, + l_mmq_wg_denoms_k, m_mmq_wg_denoms_k, s_mmq_wg_denoms_k, + l_mmqid_wg_denoms, m_mmqid_wg_denoms, s_mmqid_wg_denoms; + + uint32_t l_align, m_align, s_align; + if (device->coopmat2) { + // spec constants and tile sizes for non-quant matmul/matmul_id + l_warptile = { 256, 128, 256, 64 }; + m_warptile = { 256, 128, 128, 64 }; + s_warptile = { 128, 64, 64, 64 }; + l_wg_denoms = {128, 256, 1 }; + m_wg_denoms = {128, 128, 1 }; + s_wg_denoms = { 64, 64, 1 }; + + // spec constants and tile sizes for quant matmul (non-Qi_K) + l_warptile_mmq = { 256, 128, 256, 64 }; + m_warptile_mmq = { 256, 128, 128, 64 }; + s_warptile_mmq = { 256, 128, 128, 64 }; + l_mmq_wg_denoms = { 128, 256, 1 }; + m_mmq_wg_denoms = { 128, 128, 1 }; + s_mmq_wg_denoms = { 128, 128, 1 }; + + // spec constants and tile sizes for quant matmul (Qi_K) + l_warptile_mmq_k = { 256, 128, 512, 16 }; + m_warptile_mmq_k = { 256, 128, 256, 16 }; + s_warptile_mmq_k = { 256, 32, 128, 64 }; + l_mmq_wg_denoms_k = { 128, 512, 1 }; + m_mmq_wg_denoms_k = { 128, 256, 1 }; + s_mmq_wg_denoms_k = { 32, 128, 1 }; + + // spec constants and tile sizes for quant matmul_id + l_warptile_mmqid = { 256, 128, 128, 16 }; + m_warptile_mmqid = { 256, 128, 64, 16 }; + s_warptile_mmqid = { 256, 64, 64, 16 }; + l_mmqid_wg_denoms = { 128, 128, 1 }; + m_mmqid_wg_denoms = { 128, 64, 1 }; + s_mmqid_wg_denoms = { 64, 64, 1 }; + + l_align = 128; + m_align = 64; + s_align = 32; + } else { + // Matrix cores require different warp group sizes + const uint32_t tm_l = device->coopmat_support ? device->coopmat_m : 4; + const uint32_t tm_m = device->coopmat_support ? device->coopmat_m : 4; + const uint32_t tm_s = device->coopmat_support ? device->coopmat_m : 2; + const uint32_t tn_l = device->coopmat_support ? device->coopmat_n : 4; + const uint32_t tn_m = device->coopmat_support ? device->coopmat_n : 2; + const uint32_t tn_s = device->coopmat_support ? device->coopmat_n : 2; + const uint32_t tk_l = device->coopmat_support ? device->coopmat_k : 1; + const uint32_t tk_m = device->coopmat_support ? device->coopmat_k : 1; + const uint32_t tk_s = device->coopmat_support ? device->coopmat_k : 1; + + l_warptile = { 128, 128, 128, 16, device->subgroup_size * 2, 64, 2, tm_l, tn_l, tk_l, device->subgroup_size }; + m_warptile = { 128, 64, 64, 16, device->subgroup_size, 32, 2, tm_m, tn_m, tk_m, device->subgroup_size }; + s_warptile = { subgroup_size_16, 32, 32, 16, 32, 32, 2, tm_s, tn_s, tk_s, device->subgroup_size }; + + l_warptile_mmq = { 128, 128, 128, 32, device->subgroup_size * 2, 64, 2, tm_l, tn_l, tk_l, device->subgroup_size }; + m_warptile_mmq = { 128, 64, 64, 32, device->subgroup_size, 32, 2, tm_m, tn_m, tk_m, device->subgroup_size }; + s_warptile_mmq = { subgroup_size_32, 32, 32, 32, 32, 32, 2, tm_s, tn_s, tk_s, device->subgroup_size }; + + l_mmq_wg_denoms = l_wg_denoms = {128, 128, 1 }; + m_mmq_wg_denoms = m_wg_denoms = { 64, 64, 1 }; + s_mmq_wg_denoms = s_wg_denoms = { 32, 32, 1 }; + l_align = 128; + m_align = 64; + s_align = 32; + + // Fallback to smaller sizes if there's not enough shared memory. Given the current shaders + // and tile sizes, this should handle 16KB, 32KB, and 48KB+. + // This logic doesn't explicitly account for the 12KB row_ids in the mul_mat_mat_id shaders. + // But the numbers happen to work out for 32KB shared memory size that when using the medium + // size there's enough room for everything, and we assert for this. + uint32_t shmem_needed = (l_warptile[1] + l_warptile[2]) * (l_warptile[3] + 1) * sizeof(float); + if (shmem_needed > device->properties.limits.maxComputeSharedMemorySize) { + l_warptile = m_warptile; + l_wg_denoms = m_wg_denoms; + shmem_needed = (l_warptile[1] + l_warptile[2]) * (l_warptile[3] + 1) * sizeof(float); + GGML_ASSERT(shmem_needed <= device->properties.limits.maxComputeSharedMemorySize); + } + if (device->properties.limits.maxComputeSharedMemorySize >= 32768) { + // assert mul_mat_mat_id shaders will fit. + GGML_ASSERT(shmem_needed + 3072*4 <= device->properties.limits.maxComputeSharedMemorySize); + } + + shmem_needed = (l_warptile_mmq[1] + l_warptile_mmq[2]) * (l_warptile_mmq[3] + 1) * sizeof(float); + if (shmem_needed > device->properties.limits.maxComputeSharedMemorySize) { + if (device->properties.limits.maxComputeSharedMemorySize == 32768) { + l_warptile_mmq = m_warptile_mmq; + l_mmq_wg_denoms = m_mmq_wg_denoms; + } else { + l_warptile_mmq = s_warptile_mmq; + l_mmq_wg_denoms = s_mmq_wg_denoms; + } + shmem_needed = (l_warptile_mmq[1] + l_warptile_mmq[2]) * (l_warptile_mmq[3] + 1) * sizeof(float); + GGML_ASSERT(shmem_needed <= device->properties.limits.maxComputeSharedMemorySize); + } + if (device->properties.limits.maxComputeSharedMemorySize >= 32768) { + // assert mul_mat_mat_id shaders will fit. + GGML_ASSERT(shmem_needed + 3072*4 <= device->properties.limits.maxComputeSharedMemorySize); + } + // Disable medium and large matrix multiplication if not enough shared memory is available + // Check mmq warptiles as the largest configuration + // Throw an error if not enough for any matrix multiplication is available + if (!ggml_vk_matmul_shmem_support(device, s_warptile_mmq, false)) { + std::cerr << "ggml_vulkan: Error: Shared memory size too small for matrix multiplication." << std::endl; + throw std::runtime_error("Shared memory size too small for matrix multiplication."); + } else if (!ggml_vk_matmul_shmem_support(device, m_warptile_mmq, false)) { + device->mul_mat_m = false; + device->mul_mat_l = false; + } else if (!ggml_vk_matmul_shmem_support(device, l_warptile_mmq, false)) { + device->mul_mat_l = false; + } + + // Disable mul_mat_id if not enough shared memory is available + if (!ggml_vk_matmul_shmem_support(device, s_warptile_mmq, true)) { + device->mul_mat_id_s = false; + device->mul_mat_id_m = false; + device->mul_mat_id_l = false; + } else if (!ggml_vk_matmul_shmem_support(device, m_warptile_mmq, true)) { + device->mul_mat_id_m = false; + device->mul_mat_id_l = false; + } else if (!ggml_vk_matmul_shmem_support(device, l_warptile_mmq, true)) { + device->mul_mat_id_l = false; + } + } + + device->pipeline_matmul_f32 = std::make_shared(); + device->pipeline_matmul_f32_f16 = std::make_shared(); + + device->pipeline_matmul_id_f32 = std::make_shared(); + + std::vector> compiles; + auto const &ggml_vk_create_pipeline = [&](vk_device& device, vk_pipeline& pipeline, const std::string &name, size_t spv_size, const void* spv_data, const std::string &entrypoint, + uint32_t parameter_count, uint32_t push_constant_size, std::array wg_denoms, const std::vector& specialization_constants, + uint32_t align, bool disable_robustness = false, bool require_full_subgroups = false, uint32_t required_subgroup_size = 0) { + { + // wait until fewer than N compiles are in progress + uint32_t N = std::max(1u, std::thread::hardware_concurrency()); + std::unique_lock guard(compile_count_mutex); + while (compile_count >= N) { + compile_count_cond.wait(guard); + } + compile_count++; + } + compiles.push_back(std::async(ggml_vk_create_pipeline_func, std::ref(device), std::ref(pipeline), name, spv_size, spv_data, entrypoint, + parameter_count, push_constant_size, wg_denoms, specialization_constants, align, disable_robustness, require_full_subgroups, required_subgroup_size)); + }; + +#if defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + if (device->coopmat2) { + + auto const &fa_wg_denoms = [&](uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) -> std::array { + return {fa_rows_cols(D, clamp, type, small_rows)[0], 1, 1}; + }; + + auto const &fa_spec_constants = [&](uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) -> std::vector { + // For large number of rows, 128 invocations seems to work best. + // For small number of rows (e.g. N==1), 256 works better. But matrix granularity for 256 is 32, so we + // can't use 256 for D==80. + uint32_t wg_size = (small_rows && (D % 32) == 0) ? 256 : 128; + auto rows_cols = fa_rows_cols(D, clamp, type, small_rows); + return {wg_size, rows_cols[0], rows_cols[1], (D), clamp}; + }; + +#define CREATE_FA2(TYPE, NAMELC, D) \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][0][0], "flash_attn_f32_f16_D" #D "_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,false), fa_spec_constants(D,1,TYPE,false), 1); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][0][1], "flash_attn_f32_f16_D" #D "_aligned_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,false), fa_spec_constants(D,0,TYPE,false), fa_rows_cols(D,0,TYPE,false)[1]); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][0][0], "flash_attn_f32_f16_D" #D "_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,false), fa_spec_constants(D,1,TYPE,false), 1); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][0][1], "flash_attn_f32_f16_D" #D "_aligned_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,false), fa_spec_constants(D,0,TYPE,false), fa_rows_cols(D,0,TYPE,false)[1]); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][1][0], "flash_attn_f32_f16_D" #D "_f16acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,true), fa_spec_constants(D,1,TYPE,true), 1); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][1][1], "flash_attn_f32_f16_D" #D "_aligned_f16acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,true), fa_spec_constants(D,0,TYPE,true), fa_rows_cols(D,0,TYPE,true)[1]); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][1][0], "flash_attn_f32_f16_D" #D "_f32acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,true), fa_spec_constants(D,1,TYPE,true), 1); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][1][1], "flash_attn_f32_f16_D" #D "_aligned_f32acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,true), fa_spec_constants(D,0,TYPE,true), fa_rows_cols(D,0,TYPE,true)[1]); \ + +#define CREATE_FA(TYPE, NAMELC) \ + CREATE_FA2(TYPE, NAMELC, 64) \ + CREATE_FA2(TYPE, NAMELC, 80) \ + CREATE_FA2(TYPE, NAMELC, 96) \ + CREATE_FA2(TYPE, NAMELC, 112) \ + CREATE_FA2(TYPE, NAMELC, 128) \ + CREATE_FA2(TYPE, NAMELC, 256) + + CREATE_FA(GGML_TYPE_F16, f16) + CREATE_FA(GGML_TYPE_Q4_0, q4_0) + CREATE_FA(GGML_TYPE_Q4_1, q4_1) + CREATE_FA(GGML_TYPE_Q5_0, q5_0) + CREATE_FA(GGML_TYPE_Q5_1, q5_1) + CREATE_FA(GGML_TYPE_Q8_0, q8_0) + // K dequants currently disabled because D dimension is rounded up to 256 and runs inefficiently + //CREATE_FA(GGML_TYPE_Q2_K, q2_k) + //CREATE_FA(GGML_TYPE_Q3_K, q3_k) + //CREATE_FA(GGML_TYPE_Q4_K, q4_k) + //CREATE_FA(GGML_TYPE_Q5_K, q5_k) + //CREATE_FA(GGML_TYPE_Q6_K, q6_k) + CREATE_FA(GGML_TYPE_IQ4_NL, iq4_nl) +#undef CREATE_FA + + // Create 6 variants, {s,m,l}x{unaligned,aligned} +#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ + + // Create 2 variants, {f16,f32} accumulator +#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + + CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3) + + CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + + CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) + CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) + CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) +#undef CREATE_MM +#undef CREATE_MM2 + } else +#endif // defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + if (device->coopmat_support) { + // Create 6 variants, {s,m,l}x{unaligned,aligned} +#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1, false, true); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1, false, true); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1, false, true); \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align, false, true); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align, false, true); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align, false, true); \ + + // Create 2 variants, {f16,f32} accumulator +#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->coopmat_acc_f16_support) { \ + CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + } \ + if (device->coopmat_acc_f32_support) { \ + CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + } \ + + CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + + if (device->coopmat_acc_f16_support) { + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + } else { + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + } + + // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. + if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { + CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + + if (device->coopmat_acc_f16_support) { + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + } else { + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + } + } +#undef CREATE_MM2 +#undef CREATE_MM + } else if (device->fp16) { + // Create 6 variants, {s,m,l}x{unaligned,aligned} +#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ + + // Create 2 variants, {f16,f32} accumulator +#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + + CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. + if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { + CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + } +#undef CREATE_MM2 +#undef CREATE_MM + } else { + // Create 6 variants, {s,m,l}x{unaligned,aligned} +#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ + + CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_matmul_f16.f32acc, matmul_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_matmul_f16_f32.f32acc, matmul_f16_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f32acc, matmul_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f32acc, matmul_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f32acc, matmul_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f32acc, matmul_q5_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f32acc, matmul_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f32acc, matmul_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f32acc, matmul_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f32acc, matmul_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f32acc, matmul_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f32acc, matmul_q6_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f32acc, matmul_iq4_nl_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. + if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { + CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM(pipeline_matmul_id_f16.f32acc, matmul_id_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM(pipeline_matmul_id_f16_f32.f32acc, matmul_id_f16_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f32acc, matmul_id_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f32acc, matmul_id_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f32acc, matmul_id_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f32acc, matmul_id_q5_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f32acc, matmul_id_q6_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f32acc, matmul_id_iq4_nl_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + } +#undef CREATE_MM + } + + // mul mat vec + + // the number of rows computed per shader depends on GPU model and quant + uint32_t rm_stdq = 1; + uint32_t rm_kq = 2; + if (device->vendor_id == VK_VENDOR_ID_AMD) { + if (device->subgroup_min_size == 64 && device->subgroup_max_size == 64) { // GCN + rm_stdq = 2; + rm_kq = 4; + } + } else if (device->vendor_id == VK_VENDOR_ID_INTEL) + rm_stdq = 2; + + for (uint32_t i = 0; i < mul_mat_vec_max_cols; ++i) { + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_F32 ][i], "mul_mat_vec_f32_f32_f32_"+std::to_string(i+1), mul_mat_vec_f32_f32_f32_len, mul_mat_vec_f32_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f32_f32_"+std::to_string(i+1), mul_mat_vec_f16_f32_f32_len, mul_mat_vec_f16_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_0_f32_f32_len, mul_mat_vec_q4_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_1_f32_f32_len, mul_mat_vec_q4_1_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_0_f32_f32_len, mul_mat_vec_q5_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_1_f32_f32_len, mul_mat_vec_q5_1_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q8_0_f32_f32_len, mul_mat_vec_q8_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q2_k_f32_f32_len, mul_mat_vec_q2_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q3_k_f32_f32_len, mul_mat_vec_q3_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_k_f32_f32_len, mul_mat_vec_q4_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_k_f32_f32_len, mul_mat_vec_q5_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q6_k_f32_f32_len, mul_mat_vec_q6_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f32_f32_"+std::to_string(i+1), mul_mat_vec_iq4_nl_f32_f32_len, mul_mat_vec_iq4_nl_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq, i+1}, 1, true); + + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_F32 ][i], "mul_mat_vec_f32_f16_f32_"+std::to_string(i+1), mul_mat_vec_f32_f16_f32_len, mul_mat_vec_f32_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f16_f32_"+std::to_string(i+1), mul_mat_vec_f16_f16_f32_len, mul_mat_vec_f16_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_0_f16_f32_len, mul_mat_vec_q4_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_1_f16_f32_len, mul_mat_vec_q4_1_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_0_f16_f32_len, mul_mat_vec_q5_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_1_f16_f32_len, mul_mat_vec_q5_1_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q8_0_f16_f32_len, mul_mat_vec_q8_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q2_k_f16_f32_len, mul_mat_vec_q2_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q3_k_f16_f32_len, mul_mat_vec_q3_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_k_f16_f32_len, mul_mat_vec_q4_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_k_f16_f32_len, mul_mat_vec_q5_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q6_k_f16_f32_len, mul_mat_vec_q6_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f16_f32_"+std::to_string(i+1), mul_mat_vec_iq4_nl_f16_f32_len, mul_mat_vec_iq4_nl_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq, i+1}, 1, true); + } + + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_F32 ], "mul_mat_vec_id_f32_f32", mul_mat_vec_id_f32_f32_len, mul_mat_vec_id_f32_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {device->subgroup_size, 2}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_F16 ], "mul_mat_vec_id_f16_f32", mul_mat_vec_id_f16_f32_len, mul_mat_vec_id_f16_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {device->subgroup_size, 2}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_0], "mul_mat_vec_id_q4_0_f32", mul_mat_vec_id_q4_0_f32_len, mul_mat_vec_id_q4_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_1], "mul_mat_vec_id_q4_1_f32", mul_mat_vec_id_q4_1_f32_len, mul_mat_vec_id_q4_1_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_0], "mul_mat_vec_id_q5_0_f32", mul_mat_vec_id_q5_0_f32_len, mul_mat_vec_id_q5_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_1], "mul_mat_vec_id_q5_1_f32", mul_mat_vec_id_q5_1_f32_len, mul_mat_vec_id_q5_1_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q8_0], "mul_mat_vec_id_q8_0_f32", mul_mat_vec_id_q8_0_f32_len, mul_mat_vec_id_q8_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q2_K], "mul_mat_vec_id_q2_k_f32", mul_mat_vec_id_q2_k_f32_len, mul_mat_vec_id_q2_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q3_K], "mul_mat_vec_id_q3_k_f32", mul_mat_vec_id_q3_k_f32_len, mul_mat_vec_id_q3_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_K], "mul_mat_vec_id_q4_k_f32", mul_mat_vec_id_q4_k_f32_len, mul_mat_vec_id_q4_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_K], "mul_mat_vec_id_q5_k_f32", mul_mat_vec_id_q5_k_f32_len, mul_mat_vec_id_q5_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q6_K], "mul_mat_vec_id_q6_k_f32", mul_mat_vec_id_q6_k_f32_len, mul_mat_vec_id_q6_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_IQ4_NL], "mul_mat_vec_id_iq4_nl_f32", mul_mat_vec_id_iq4_nl_f32_len, mul_mat_vec_id_iq4_nl_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq}, 1, true); + + // dequant shaders + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_F32 ], "f32_to_f16", dequant_f32_len, dequant_f32_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_0], "dequant_q4_0", dequant_q4_0_len, dequant_q4_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_1], "dequant_q4_1", dequant_q4_1_len, dequant_q4_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_1], "dequant_q5_1", dequant_q5_1_len, dequant_q5_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_K], "dequant_q4_k", dequant_q4_k_len, dequant_q4_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_K], "dequant_q5_k", dequant_q5_k_len, dequant_q5_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q6_K], "dequant_q6_k", dequant_q6_k_len, dequant_q6_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ4_NL], "dequant_iq4_nl", dequant_iq4_nl_len, dequant_iq4_nl_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + + // get_rows + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_F32 ], "get_rows_f32", get_rows_f32_len, get_rows_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_F16 ], "get_rows_f16", get_rows_f16_len, get_rows_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_0], "get_rows_q4_0", get_rows_q4_0_len, get_rows_q4_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_1], "get_rows_q4_1", get_rows_q4_1_len, get_rows_q4_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_0], "get_rows_q5_0", get_rows_q5_0_len, get_rows_q5_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_1], "get_rows_q5_1", get_rows_q5_1_len, get_rows_q5_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q8_0], "get_rows_q8_0", get_rows_q8_0_len, get_rows_q8_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_IQ4_NL], "get_rows_iq4_nl", get_rows_iq4_nl_len, get_rows_iq4_nl_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_F32 ], "get_rows_f32_f32", get_rows_f32_f32_len, get_rows_f32_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_F16 ], "get_rows_f16_f32", get_rows_f16_f32_len, get_rows_f16_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_0], "get_rows_q4_0_f32", get_rows_q4_0_f32_len, get_rows_q4_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_1], "get_rows_q4_1_f32", get_rows_q4_1_f32_len, get_rows_q4_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_0], "get_rows_q5_0_f32", get_rows_q5_0_f32_len, get_rows_q5_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_1], "get_rows_q5_1_f32", get_rows_q5_1_f32_len, get_rows_q5_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q8_0], "get_rows_q8_0_f32", get_rows_q8_0_f32_len, get_rows_q8_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_IQ4_NL], "get_rows_iq4_nl_f32", get_rows_iq4_nl_f32_len, get_rows_iq4_nl_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_matmul_split_k_reduce, "split_k_reduce", split_k_reduce_len, split_k_reduce_data, "main", 2, 2 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_mul_mat_vec_p021_f16_f32, "mul_mat_vec_p021_f16_f32", mul_mat_vec_p021_f16_f32_len, mul_mat_vec_p021_f16_f32_data, "main", 3, 6 * sizeof(uint32_t), {1, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_mul_mat_vec_nc_f16_f32, "mul_mat_vec_nc_f16_f32", mul_mat_vec_nc_f16_f32_len, mul_mat_vec_nc_f16_f32_data, "main", 3, 7 * sizeof(uint32_t), {1, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_norm_f32, "norm_f32", norm_f32_len, norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_group_norm_f32, "group_norm_f32", group_norm_f32_len, group_norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_f32, "rms_norm_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f32, "cpy_f32_f32", cpy_f32_f32_len, cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f16, "cpy_f32_f16", cpy_f32_f16_len, cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_f16_f16, "cpy_f16_f16", cpy_f16_f16_len, cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f32, "contig_cpy_f32_f32", contig_cpy_f32_f32_len, contig_cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f16, "contig_cpy_f32_f16", contig_cpy_f32_f16_len, contig_cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f16_f16, "contig_cpy_f16_f16", contig_cpy_f16_f16_len, contig_cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_add_f32, "add_f32", add_f32_len, add_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); + ggml_vk_create_pipeline(device, device->pipeline_add_f32_norepeat, "add_f32_norepeat", add_f32_len, add_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_add_f16_f32_f16, "add_f16_f32_f16", add_f16_f32_f16_len, add_f16_f32_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); + ggml_vk_create_pipeline(device, device->pipeline_add_f16_f32_f16_norepeat, "add_f16_f32_f16_norepeat", add_f16_f32_f16_len, add_f16_f32_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_acc_f32, "acc_f32", acc_f32_len, acc_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_mul_f32, "mul_f32", mul_f32_len, mul_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); + ggml_vk_create_pipeline(device, device->pipeline_mul_f32_norepeat, "mul_f32_norepeat", mul_f32_len, mul_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_div_f32, "div_f32", div_f32_len, div_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); + ggml_vk_create_pipeline(device, device->pipeline_div_f32_norepeat, "div_f32_norepeat", div_f32_len, div_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_concat_f32, "concat_f32", concat_f32_len, concat_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_concat_f16, "concat_f16", concat_f16_len, concat_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_concat_i32, "concat_i32", concat_i32_len, concat_i32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_upscale_f32, "upscale_f32", upscale_f32_len, upscale_f32_data, "main", 2, sizeof(vk_op_upscale_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_scale_f32, "scale_f32", scale_f32_len, scale_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_sqr_f32, "sqr_f32", sqr_f32_len, sqr_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_sin_f32, "sin_f32", sin_f32_len, sin_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cos_f32, "cos_f32", cos_f32_len, cos_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_clamp_f32, "clamp_f32", clamp_f32_len, clamp_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_pad_f32, "pad_f32", pad_f32_len, pad_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_repeat_f32, "repeat_f32", repeat_f32_len, repeat_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_gelu_f32, "gelu_f32", gelu_f32_len, gelu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_gelu_quick_f32, "gelu_quick_f32", gelu_quick_f32_len, gelu_quick_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_silu_f32, "silu_f32", silu_f32_len, silu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_relu_f32, "relu_f32", relu_f32_len, relu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_leaky_relu_f32, "leaky_relu_f32", leaky_relu_f32_len, leaky_relu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_tanh_f32, "tanh_f32", tanh_f32_len, tanh_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_diag_mask_inf_f32, "diag_mask_inf_f32", diag_mask_inf_f32_len, diag_mask_inf_f32_data, "main", 2, sizeof(vk_op_diag_mask_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32, "soft_max_f32", soft_max_f32_len, soft_max_f32_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); + ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_wg512, "soft_max_f32_wg512", soft_max_f32_len, soft_max_f32_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 512 }, 1); + ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_f16, "soft_max_f32_f16", soft_max_f32_f16_len, soft_max_f32_f16_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); + ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_f16_wg512, "soft_max_f32_f16_wg512", soft_max_f32_f16_len, soft_max_f32_f16_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 512 }, 1); + + ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f32, "rope_norm_f32", rope_norm_f32_len, rope_norm_f32_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f32, "rope_neox_f32", rope_neox_f32_len, rope_neox_f32_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + + if (device->float_controls_rte_fp16) { + ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f16, "rope_norm_f16", rope_norm_f16_rte_len, rope_norm_f16_rte_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f16, "rope_neox_f16", rope_neox_f16_rte_len, rope_neox_f16_rte_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + } else { + ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f16, "rope_norm_f16", rope_norm_f16_len, rope_norm_f16_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f16, "rope_neox_f16", rope_neox_f16_len, rope_neox_f16_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + } + + ggml_vk_create_pipeline(device, device->pipeline_argsort_f32, "argsort_f32", argsort_f32_len, argsort_f32_data, "main", 2, sizeof(vk_op_argsort_push_constants), {1024, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); + + ggml_vk_create_pipeline(device, device->pipeline_im2col_f32, "im2col_f32", im2col_f32_len, im2col_f32_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); + if (device->float_controls_rte_fp16) { + ggml_vk_create_pipeline(device, device->pipeline_im2col_f32_f16, "im2col_f32_f16", im2col_f32_f16_rte_len, im2col_f32_f16_rte_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); + } else { + ggml_vk_create_pipeline(device, device->pipeline_im2col_f32_f16, "im2col_f32_f16", im2col_f32_f16_len, im2col_f32_f16_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); + } + + ggml_vk_create_pipeline(device, device->pipeline_timestep_embedding_f32, "timestep_embedding_f32", timestep_embedding_f32_len, timestep_embedding_f32_data, "main", 2, sizeof(vk_op_timestep_embedding_push_constants), {256, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_pool2d_f32, "pool2d_f32", pool2d_f32_len, pool2d_f32_data, "main", 2, sizeof(vk_op_pool2d_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv6_f32, "rwkv_wkv6_f32", rwkv_wkv6_f32_len, rwkv_wkv6_f32_data, "main", 7, sizeof(vk_op_rwkv_wkv6_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); + + for (auto &c : compiles) { + c.wait(); + } + std::cerr << "Done!" << std::endl; +} + +static bool ggml_vk_khr_cooperative_matrix_support(const vk::PhysicalDeviceProperties& props, const vk::PhysicalDeviceDriverProperties& driver_props); + +static vk_device ggml_vk_get_device(size_t idx) { + VK_LOG_DEBUG("ggml_vk_get_device(" << idx << ")"); + + if (vk_instance.devices[idx] == nullptr) { + VK_LOG_DEBUG("Initializing new vk_device"); + vk_device device = std::make_shared(); + vk_instance.devices[idx] = device; + +#ifdef GGML_VULKAN_MEMORY_DEBUG + device->memory_logger = std::unique_ptr(new vk_memory_logger()); +#endif +#ifdef GGML_VULKAN_PERF + device->perf_logger = std::unique_ptr(new vk_perf_logger()); +#endif + + size_t dev_num = vk_instance.device_indices[idx]; + + std::vector physical_devices = vk_instance.instance.enumeratePhysicalDevices(); + + if (dev_num >= physical_devices.size()) { + std::cerr << "ggml_vulkan: Device with index " << dev_num << " does not exist." << std::endl; + throw std::runtime_error("Device not found"); + } + + device->physical_device = physical_devices[dev_num]; + const std::vector ext_props = device->physical_device.enumerateDeviceExtensionProperties(); + + bool fp16_storage = false; + bool fp16_compute = false; + bool maintenance4_support = false; + bool sm_builtins = false; + bool amd_shader_core_properties2 = false; + bool pipeline_robustness = false; + bool coopmat2_support = false; + device->coopmat_support = false; + + // Check if maintenance4 is supported + for (const auto& properties : ext_props) { + if (strcmp("VK_KHR_maintenance4", properties.extensionName) == 0) { + maintenance4_support = true; + } else if (strcmp("VK_KHR_16bit_storage", properties.extensionName) == 0) { + fp16_storage = true; + } else if (strcmp("VK_KHR_shader_float16_int8", properties.extensionName) == 0) { + fp16_compute = true; + } else if (strcmp("VK_NV_shader_sm_builtins", properties.extensionName) == 0) { + sm_builtins = true; + } else if (strcmp("VK_AMD_shader_core_properties2", properties.extensionName) == 0) { + amd_shader_core_properties2 = true; + } else if (strcmp("VK_EXT_pipeline_robustness", properties.extensionName) == 0) { + pipeline_robustness = true; + } else if (strcmp("VK_EXT_subgroup_size_control", properties.extensionName) == 0) { + device->subgroup_size_control = true; + } else if (strcmp("VK_KHR_cooperative_matrix", properties.extensionName) == 0 && + !getenv("GGML_VK_DISABLE_COOPMAT")) { + device->coopmat_support = true; + device->coopmat_m = 0; + device->coopmat_n = 0; + device->coopmat_k = 0; + } else if (strcmp("VK_NV_cooperative_matrix2", properties.extensionName) == 0 && + !getenv("GGML_VK_DISABLE_COOPMAT2")) { + coopmat2_support = true; + } + } + + vk::PhysicalDeviceProperties2 props2; + vk::PhysicalDeviceMaintenance3Properties props3; + vk::PhysicalDeviceMaintenance4Properties props4; + vk::PhysicalDeviceSubgroupProperties subgroup_props; + vk::PhysicalDeviceDriverProperties driver_props; + vk::PhysicalDeviceShaderSMBuiltinsPropertiesNV sm_props; + vk::PhysicalDeviceShaderCoreProperties2AMD amd_shader_core_properties2_props; + vk::PhysicalDeviceVulkan12Properties vk12_props; + vk::PhysicalDeviceSubgroupSizeControlPropertiesEXT subgroup_size_control_props; + + props2.pNext = &props3; + props3.pNext = &subgroup_props; + subgroup_props.pNext = &driver_props; + driver_props.pNext = &vk12_props; + + VkBaseOutStructure * last_struct = (VkBaseOutStructure *)&vk12_props; + + if (maintenance4_support) { + last_struct->pNext = (VkBaseOutStructure *)&props4; + last_struct = (VkBaseOutStructure *)&props4; + } + if (sm_builtins) { + last_struct->pNext = (VkBaseOutStructure *)&sm_props; + last_struct = (VkBaseOutStructure *)&sm_props; + } + if (amd_shader_core_properties2) { + last_struct->pNext = (VkBaseOutStructure *)&amd_shader_core_properties2_props; + last_struct = (VkBaseOutStructure *)&amd_shader_core_properties2_props; + } + if (device->subgroup_size_control) { + last_struct->pNext = (VkBaseOutStructure *)&subgroup_size_control_props; + last_struct = (VkBaseOutStructure *)&subgroup_size_control_props; + } + +#if defined(VK_NV_cooperative_matrix2) + vk::PhysicalDeviceCooperativeMatrix2PropertiesNV coopmat2_props; + if (coopmat2_support) { + last_struct->pNext = (VkBaseOutStructure *)&coopmat2_props; + last_struct = (VkBaseOutStructure *)&coopmat2_props; + } +#endif + + device->physical_device.getProperties2(&props2); + device->properties = props2.properties; + + const char* GGML_VK_FORCE_MAX_ALLOCATION_SIZE = getenv("GGML_VK_FORCE_MAX_ALLOCATION_SIZE"); + + if (GGML_VK_FORCE_MAX_ALLOCATION_SIZE != nullptr) { + device->max_memory_allocation_size = std::stoul(GGML_VK_FORCE_MAX_ALLOCATION_SIZE); + } else if (maintenance4_support) { + device->max_memory_allocation_size = std::min(props3.maxMemoryAllocationSize, props4.maxBufferSize); + } else { + device->max_memory_allocation_size = props3.maxMemoryAllocationSize; + } + + device->vendor_id = device->properties.vendorID; + device->subgroup_size = subgroup_props.subgroupSize; + device->uma = device->properties.deviceType == vk::PhysicalDeviceType::eIntegratedGpu; + if (sm_builtins) { + device->shader_core_count = sm_props.shaderSMCount; + } else if (amd_shader_core_properties2) { + device->shader_core_count = amd_shader_core_properties2_props.activeComputeUnitCount; + } else { + device->shader_core_count = 0; + } + device->float_controls_rte_fp16 = vk12_props.shaderRoundingModeRTEFloat16; + + const bool force_disable_f16 = getenv("GGML_VK_DISABLE_F16") != nullptr; + + device->fp16 = !force_disable_f16 && fp16_storage && fp16_compute; + + if (!ggml_vk_khr_cooperative_matrix_support(device->properties, driver_props)) { + device->coopmat_support = false; + } + + std::vector queue_family_props = device->physical_device.getQueueFamilyProperties(); + + // Try to find a non-graphics compute queue and transfer-focused queues + const uint32_t compute_queue_family_index = ggml_vk_find_queue_family_index(queue_family_props, vk::QueueFlagBits::eCompute, vk::QueueFlagBits::eGraphics, -1, 1); + const uint32_t transfer_queue_family_index = ggml_vk_find_queue_family_index(queue_family_props, vk::QueueFlagBits::eTransfer, vk::QueueFlagBits::eCompute | vk::QueueFlagBits::eGraphics, compute_queue_family_index, 1); + + const float priorities[] = { 1.0f, 1.0f }; + device->single_queue = compute_queue_family_index == transfer_queue_family_index && queue_family_props[compute_queue_family_index].queueCount == 1; + + std::vector device_queue_create_infos; + if (compute_queue_family_index != transfer_queue_family_index) { + device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities}); + device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), transfer_queue_family_index, 1, priorities + 1}); + } else if(!device->single_queue) { + device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 2, priorities}); + } else { + device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities}); + } + vk::DeviceCreateInfo device_create_info; + std::vector device_extensions; + vk::PhysicalDeviceFeatures device_features = device->physical_device.getFeatures(); + + VkPhysicalDeviceFeatures2 device_features2; + device_features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2; + device_features2.pNext = nullptr; + device_features2.features = (VkPhysicalDeviceFeatures)device_features; + + VkPhysicalDeviceVulkan11Features vk11_features; + vk11_features.pNext = nullptr; + vk11_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES; + device_features2.pNext = &vk11_features; + + VkPhysicalDeviceVulkan12Features vk12_features; + vk12_features.pNext = nullptr; + vk12_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES; + vk11_features.pNext = &vk12_features; + + last_struct = (VkBaseOutStructure *)&vk12_features; + + VkPhysicalDevicePipelineRobustnessFeaturesEXT pl_robustness_features; + pl_robustness_features.pNext = nullptr; + pl_robustness_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PIPELINE_ROBUSTNESS_FEATURES_EXT; + pl_robustness_features.pipelineRobustness = VK_FALSE; + + if (pipeline_robustness) { + last_struct->pNext = (VkBaseOutStructure *)&pl_robustness_features; + last_struct = (VkBaseOutStructure *)&pl_robustness_features; + device_extensions.push_back("VK_EXT_pipeline_robustness"); + } + + VkPhysicalDeviceSubgroupSizeControlFeaturesEXT subgroup_size_control_features; + subgroup_size_control_features.pNext = nullptr; + subgroup_size_control_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES_EXT; + subgroup_size_control_features.computeFullSubgroups = false; + subgroup_size_control_features.subgroupSizeControl = false; + + if (device->subgroup_size_control) { + last_struct->pNext = (VkBaseOutStructure *)&subgroup_size_control_features; + last_struct = (VkBaseOutStructure *)&subgroup_size_control_features; + } + + VkPhysicalDeviceCooperativeMatrixFeaturesKHR coopmat_features; + coopmat_features.pNext = nullptr; + coopmat_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR; + coopmat_features.cooperativeMatrix = VK_FALSE; + + if (device->coopmat_support) { + last_struct->pNext = (VkBaseOutStructure *)&coopmat_features; + last_struct = (VkBaseOutStructure *)&coopmat_features; + } + +#if defined(VK_NV_cooperative_matrix2) + VkPhysicalDeviceCooperativeMatrix2FeaturesNV coopmat2_features {}; + coopmat2_features.pNext = nullptr; + coopmat2_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_2_FEATURES_NV; + if (coopmat2_support) { + last_struct->pNext = (VkBaseOutStructure *)&coopmat2_features; + last_struct = (VkBaseOutStructure *)&coopmat2_features; + device_extensions.push_back("VK_NV_cooperative_matrix2"); + } +#endif + + vkGetPhysicalDeviceFeatures2(device->physical_device, &device_features2); + + device->fp16 = device->fp16 && vk12_features.shaderFloat16; + + device->pipeline_robustness = pl_robustness_features.pipelineRobustness; + + if (device->subgroup_size_control) { + device->subgroup_min_size = subgroup_size_control_props.minSubgroupSize; + device->subgroup_max_size = subgroup_size_control_props.maxSubgroupSize; + } + + device->subgroup_size_control = device->subgroup_size_control && + (subgroup_size_control_props.requiredSubgroupSizeStages & vk::ShaderStageFlagBits::eCompute) && + subgroup_size_control_features.subgroupSizeControl; + + if (device->subgroup_size_control) { + device->subgroup_require_full_support = subgroup_size_control_features.computeFullSubgroups; + device_extensions.push_back("VK_EXT_subgroup_size_control"); + } + + device->coopmat_support = device->coopmat_support && coopmat_features.cooperativeMatrix; + + if (coopmat2_support) { +#if defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + if (coopmat2_features.cooperativeMatrixWorkgroupScope && + coopmat2_features.cooperativeMatrixFlexibleDimensions && + coopmat2_features.cooperativeMatrixReductions && + coopmat2_features.cooperativeMatrixConversions && + coopmat2_features.cooperativeMatrixPerElementOperations && + coopmat2_features.cooperativeMatrixTensorAddressing && + coopmat2_features.cooperativeMatrixBlockLoads && + vk12_features.bufferDeviceAddress) { + + std::vector flexible_dimensions; + uint32_t count = 0; + + PFN_vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV + _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV = + (PFN_vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV) + vk_instance.instance.getProcAddr("vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV"); + + _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV(device->physical_device, &count, nullptr); + + VkCooperativeMatrixFlexibleDimensionsPropertiesNV empty_prop {}; + empty_prop.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_FLEXIBLE_DIMENSIONS_PROPERTIES_NV; + flexible_dimensions.resize(count, empty_prop); + + _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV(device->physical_device, &count, flexible_dimensions.data()); + + bool found_fp16_128 = false, + found_fp16_256 = false, + found_fp32_128 = false, + found_fp32_256 = false; + // need to support fp16*fp16 with fp16/fp32 accumulator, for workgroupsize 128 + // with 32x16x16 and 256 with 32x32x16. + for (auto &prop : flexible_dimensions) { + if (prop.saturatingAccumulation == VK_FALSE && + prop.scope == VK_SCOPE_WORKGROUP_KHR && + prop.AType == VK_COMPONENT_TYPE_FLOAT16_KHR && + prop.BType == VK_COMPONENT_TYPE_FLOAT16_KHR) { + + if (prop.workgroupInvocations == 128 && + prop.MGranularity <= 32 && + prop.NGranularity <= 16 && + prop.KGranularity <= 16) { + if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR && + prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) { + found_fp16_128 = true; + } + if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && + prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) { + found_fp32_128 = true; + } + } + if (prop.workgroupInvocations == 256 && + prop.MGranularity <= 32 && + prop.NGranularity <= 32 && + prop.KGranularity <= 16) { + if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR && + prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) { + found_fp16_256 = true; + } + if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && + prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) { + found_fp32_256 = true; + } + } + } + } + if (found_fp16_128 && found_fp16_256 && + found_fp32_128 && found_fp32_256 && + coopmat2_props.cooperativeMatrixFlexibleDimensionsMaxDimension >= 512) { + device->coopmat2 = true; + } + } +#endif + } + + if (!vk11_features.storageBuffer16BitAccess) { + std::cerr << "ggml_vulkan: device " << GGML_VK_NAME << idx << " does not support 16-bit storage." << std::endl; + throw std::runtime_error("Unsupported device"); + } + + device_extensions.push_back("VK_KHR_16bit_storage"); + +#ifdef GGML_VULKAN_VALIDATE + device_extensions.push_back("VK_KHR_shader_non_semantic_info"); +#endif + + if (device->fp16) { + device_extensions.push_back("VK_KHR_shader_float16_int8"); + } + + if (device->coopmat_support) { + // Query supported shapes + std::vector cm_props; + + PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR = + (PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR)vkGetInstanceProcAddr(vk_instance.instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR"); + + uint32_t cm_props_num; + + pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR(device->physical_device, &cm_props_num, nullptr); + + cm_props.resize(cm_props_num); + + for (auto& prop : cm_props) { + prop.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR; + } + + pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR(device->physical_device, &cm_props_num, cm_props.data()); + + VK_LOG_DEBUG("ggml_vulkan: Cooperative Matrix Shapes: " << cm_props.size()); + + for (auto& prop : cm_props) { + VK_LOG_DEBUG("ggml_vulkan: M: " << prop.MSize << " N: " << prop.NSize << " K: " << prop.KSize << " A: " << vk::to_string((vk::ComponentTypeKHR)prop.AType) << " B: " << vk::to_string((vk::ComponentTypeKHR)prop.BType) << " C: " << vk::to_string((vk::ComponentTypeKHR)prop.CType) << " Result: " << vk::to_string((vk::ComponentTypeKHR)prop.ResultType) << " saturatingAccumulation: " << prop.saturatingAccumulation << " scope: " << vk::to_string((vk::ScopeKHR)prop.scope)); + + if ((vk::ComponentTypeKHR)prop.AType == vk::ComponentTypeKHR::eFloat16 && + (vk::ComponentTypeKHR)prop.BType == vk::ComponentTypeKHR::eFloat16 && + (vk::ScopeKHR)prop.scope == vk::ScopeKHR::eSubgroup + ) { + if ((vk::ComponentTypeKHR)prop.CType == vk::ComponentTypeKHR::eFloat32 && + (vk::ComponentTypeKHR)prop.ResultType == vk::ComponentTypeKHR::eFloat32) { + // coopmat sizes not set yet + if (device->coopmat_m == 0) { + device->coopmat_acc_f32_support = true; + device->coopmat_m = prop.MSize; + device->coopmat_n = prop.NSize; + device->coopmat_k = prop.KSize; + } else if (device->coopmat_m == prop.MSize && device->coopmat_n == prop.NSize && device->coopmat_k == prop.KSize) { + // Only enable if shape is identical + device->coopmat_acc_f32_support = true; + } + } else if ((vk::ComponentTypeKHR)prop.CType == vk::ComponentTypeKHR::eFloat16 && + (vk::ComponentTypeKHR)prop.ResultType == vk::ComponentTypeKHR::eFloat16) { + // coopmat sizes not set yet + if (device->coopmat_m == 0) { + device->coopmat_acc_f16_support = true; + device->coopmat_m = prop.MSize; + device->coopmat_n = prop.NSize; + device->coopmat_k = prop.KSize; + } else if (device->coopmat_m == prop.MSize && device->coopmat_n == prop.NSize && device->coopmat_k == prop.KSize) { + // Only enable if shape is identical + device->coopmat_acc_f16_support = true; + } + } + } + } + + if (device->coopmat_m == 0 || !device->coopmat_acc_f32_support) { + // No suitable matmul mode found + GGML_LOG_DEBUG("ggml_vulkan: WARNING: No suitable matrix core mode found. Disabling matrix cores.\n"); + device->coopmat_support = false; + } + } + + if (device->coopmat_support) { + device_extensions.push_back("VK_KHR_cooperative_matrix"); + } + + device->name = GGML_VK_NAME + std::to_string(idx); + + device_create_info = { + vk::DeviceCreateFlags(), + device_queue_create_infos, + {}, + device_extensions + }; + device_create_info.setPNext(&device_features2); + device->device = device->physical_device.createDevice(device_create_info); + + // Queues + ggml_vk_create_queue(device, device->compute_queue, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false); + + // Shaders + // Disable matmul tile sizes early if performance low or not supported + switch (device->vendor_id) { +#ifndef GGML_VULKAN_RUN_TESTS + case VK_VENDOR_ID_AMD: + case VK_VENDOR_ID_INTEL: + device->mul_mat_l = false; + device->mul_mat_m = true; + device->mul_mat_s = true; + device->mul_mat_id_l = false; + device->mul_mat_id_m = true; + device->mul_mat_id_s = true; + break; + case VK_VENDOR_ID_APPLE: + device->mul_mat_l = false; + device->mul_mat_m = true; + device->mul_mat_s = false; + device->mul_mat_id_l = false; + device->mul_mat_id_m = true; + device->mul_mat_id_s = false; + break; +#endif + default: + device->mul_mat_l = true; + device->mul_mat_m = true; + device->mul_mat_s = true; + device->mul_mat_id_l = true; + device->mul_mat_id_m = true; + device->mul_mat_id_s = true; + break; + } + + ggml_vk_load_shaders(device); + + if (!device->single_queue) { + const uint32_t transfer_queue_index = compute_queue_family_index == transfer_queue_family_index ? 1 : 0; + ggml_vk_create_queue(device, device->transfer_queue, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true); + } else { + // TODO: Use pointer or reference to avoid copy + device->transfer_queue = device->compute_queue; + } + + device->buffer_type = { + /* .iface = */ ggml_backend_vk_buffer_type_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), idx), + /* .context = */ new ggml_backend_vk_buffer_type_context{ device->name, device }, + }; + + device->fence = device->device.createFence({}); + + device->idx = idx; + + return device; + } + + return vk_instance.devices[idx]; +} + +static void ggml_vk_print_gpu_info(size_t idx) { + GGML_ASSERT(idx < vk_instance.device_indices.size()); + size_t dev_num = vk_instance.device_indices[idx]; + VK_LOG_DEBUG("ggml_vk_print_gpu_info(" << dev_num << ")"); + GGML_ASSERT(vk_instance_initialized); + + std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); + + if (dev_num >= devices.size()) { + std::cerr << "ggml_vulkan: Device with index " << dev_num << " does not exist." << std::endl; + throw std::runtime_error("Device not found"); + } + + vk::PhysicalDevice physical_device = devices[dev_num]; + std::vector ext_props = physical_device.enumerateDeviceExtensionProperties(); + + vk::PhysicalDeviceProperties2 props2; + vk::PhysicalDeviceMaintenance3Properties props3; + vk::PhysicalDeviceSubgroupProperties subgroup_props; + vk::PhysicalDeviceDriverProperties driver_props; + props2.pNext = &props3; + props3.pNext = &subgroup_props; + subgroup_props.pNext = &driver_props; + physical_device.getProperties2(&props2); + + const size_t subgroup_size = subgroup_props.subgroupSize; + const bool uma = props2.properties.deviceType == vk::PhysicalDeviceType::eIntegratedGpu; + + bool fp16_storage = false; + bool fp16_compute = false; + bool coopmat_support = false; + bool coopmat2_support = false; + + for (auto properties : ext_props) { + if (strcmp("VK_KHR_16bit_storage", properties.extensionName) == 0) { + fp16_storage = true; + } else if (strcmp("VK_KHR_shader_float16_int8", properties.extensionName) == 0) { + fp16_compute = true; + } else if (strcmp("VK_KHR_cooperative_matrix", properties.extensionName) == 0 && + !getenv("GGML_VK_DISABLE_COOPMAT")) { + coopmat_support = true; +#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + } else if (strcmp("VK_NV_cooperative_matrix2", properties.extensionName) == 0 && + !getenv("GGML_VK_DISABLE_COOPMAT2")) { + coopmat2_support = true; +#endif + } + } + + if (!ggml_vk_khr_cooperative_matrix_support(props2.properties, driver_props)) { + coopmat_support = false; + } + + const char* GGML_VK_DISABLE_F16 = getenv("GGML_VK_DISABLE_F16"); + bool force_disable_f16 = GGML_VK_DISABLE_F16 != nullptr; + + bool fp16 = !force_disable_f16 && fp16_storage && fp16_compute; + + vk::PhysicalDeviceFeatures device_features = physical_device.getFeatures(); + + VkPhysicalDeviceFeatures2 device_features2; + device_features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2; + device_features2.pNext = nullptr; + device_features2.features = (VkPhysicalDeviceFeatures)device_features; + + VkPhysicalDeviceVulkan11Features vk11_features; + vk11_features.pNext = nullptr; + vk11_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES; + device_features2.pNext = &vk11_features; + + VkPhysicalDeviceVulkan12Features vk12_features; + vk12_features.pNext = nullptr; + vk12_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES; + vk11_features.pNext = &vk12_features; + + // Pointer to the last chain element + VkBaseOutStructure * last_struct = (VkBaseOutStructure *)&vk12_features; + + VkPhysicalDeviceCooperativeMatrixFeaturesKHR coopmat_features; + coopmat_features.pNext = nullptr; + coopmat_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR; + coopmat_features.cooperativeMatrix = VK_FALSE; + + if (coopmat_support) { + last_struct->pNext = (VkBaseOutStructure *)&coopmat_features; + last_struct = (VkBaseOutStructure *)&coopmat_features; + } + + vkGetPhysicalDeviceFeatures2(physical_device, &device_features2); + + fp16 = fp16 && vk12_features.shaderFloat16; + + coopmat_support = coopmat_support && coopmat_features.cooperativeMatrix; + + std::string matrix_cores = coopmat2_support ? "NV_coopmat2" : coopmat_support ? "KHR_coopmat" : "none"; + + std::string device_name = props2.properties.deviceName.data(); + GGML_LOG_DEBUG("ggml_vulkan: %zu = %s (%s) | uma: %d | fp16: %d | warp size: %zu | matrix cores: %s\n", + idx, device_name.c_str(), driver_props.driverName.data(), uma, fp16, subgroup_size, matrix_cores.c_str()); + + if (props2.properties.deviceType == vk::PhysicalDeviceType::eCpu) { + GGML_LOG_DEBUG("ggml_vulkan: Warning: Device type is CPU. This is probably not the device you want.\n"); + } +} + +static bool ggml_vk_instance_validation_ext_available(const std::vector& instance_extensions); +static bool ggml_vk_instance_portability_enumeration_ext_available(const std::vector& instance_extensions); + +void ggml_vk_instance_init() { + if (vk_instance_initialized) { + return; + } + VK_LOG_DEBUG("ggml_vk_instance_init()"); + + vk_instance_initialized = true; + + vk::ApplicationInfo app_info{ "ggml-vulkan", 1, nullptr, 0, VK_API_VERSION }; + + const std::vector instance_extensions = vk::enumerateInstanceExtensionProperties(); + const bool validation_ext = ggml_vk_instance_validation_ext_available(instance_extensions); +#ifdef __APPLE__ + const bool portability_enumeration_ext = ggml_vk_instance_portability_enumeration_ext_available(instance_extensions); +#endif + + std::vector layers; + + if (validation_ext) { + layers.push_back("VK_LAYER_KHRONOS_validation"); + } + std::vector extensions; + if (validation_ext) { + extensions.push_back("VK_EXT_validation_features"); + } +#ifdef __APPLE__ + if (portability_enumeration_ext) { + extensions.push_back("VK_KHR_portability_enumeration"); + } +#endif + vk::InstanceCreateInfo instance_create_info(vk::InstanceCreateFlags{}, &app_info, layers, extensions); +#ifdef __APPLE__ + if (portability_enumeration_ext) { + instance_create_info.flags |= vk::InstanceCreateFlagBits::eEnumeratePortabilityKHR; + } +#endif + + std::vector features_enable; + vk::ValidationFeaturesEXT validation_features; + + if (validation_ext) { + features_enable = { vk::ValidationFeatureEnableEXT::eBestPractices }; + validation_features = { + features_enable, + {}, + }; + validation_features.setPNext(nullptr); + instance_create_info.setPNext(&validation_features); + GGML_LOG_DEBUG("ggml_vulkan: Validation layers enabled\n"); + } + vk_instance.instance = vk::createInstance(instance_create_info); + + size_t num_available_devices = vk_instance.instance.enumeratePhysicalDevices().size(); + + // Emulate behavior of CUDA_VISIBLE_DEVICES for Vulkan + char * devices_env = getenv("GGML_VK_VISIBLE_DEVICES"); + if (devices_env != nullptr) { + std::string devices(devices_env); + std::replace(devices.begin(), devices.end(), ',', ' '); + + std::stringstream ss(devices); + size_t tmp; + while (ss >> tmp) { + if(tmp >= num_available_devices) { + std::cerr << "ggml_vulkan: Invalid device index " << tmp << " in GGML_VK_VISIBLE_DEVICES." << std::endl; + throw std::runtime_error("Invalid Vulkan device index"); + } + vk_instance.device_indices.push_back(tmp); + } + } else { + std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); + + // Make sure at least one device exists + if (devices.empty()) { + std::cerr << "ggml_vulkan: Error: No devices found." << std::endl; + GGML_ABORT("fatal error"); + } + + // Default to using all dedicated GPUs + for (size_t i = 0; i < devices.size(); i++) { + vk::PhysicalDeviceProperties2 new_props; + vk::PhysicalDeviceDriverProperties new_driver; + vk::PhysicalDeviceIDProperties new_id; + new_props.pNext = &new_driver; + new_driver.pNext = &new_id; + devices[i].getProperties2(&new_props); + + if (new_props.properties.deviceType == vk::PhysicalDeviceType::eDiscreteGpu) { + // Check if there are two physical devices corresponding to the same GPU + auto old_device = std::find_if( + vk_instance.device_indices.begin(), + vk_instance.device_indices.end(), + [&devices, &new_id](const size_t k){ + vk::PhysicalDeviceProperties2 old_props; + vk::PhysicalDeviceIDProperties old_id; + old_props.pNext = &old_id; + devices[k].getProperties2(&old_props); + return std::equal(std::begin(old_id.deviceUUID), std::end(old_id.deviceUUID), std::begin(new_id.deviceUUID)); + } + ); + if (old_device == vk_instance.device_indices.end()) { + vk_instance.device_indices.push_back(i); + } else { + // There can be two physical devices corresponding to the same GPU if there are 2 different drivers + // This can cause error when splitting layers aross the devices, need to keep only 1 + VK_LOG_DEBUG("Device " << i << " and device " << *old_device << " have the same deviceUUID"); + + vk::PhysicalDeviceProperties2 old_props; + vk::PhysicalDeviceDriverProperties old_driver; + old_props.pNext = &old_driver; + devices[*old_device].getProperties2(&old_props); + + std::map driver_priorities {}; + int old_priority = std::numeric_limits::max(); + int new_priority = std::numeric_limits::max(); + + // Check https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/VkDriverId.html for the list of driver id + // Smaller number -> higher priority + switch (old_props.properties.vendorID) { + case VK_VENDOR_ID_AMD: + driver_priorities[vk::DriverId::eMesaRadv] = 1; + driver_priorities[vk::DriverId::eAmdOpenSource] = 2; + driver_priorities[vk::DriverId::eAmdProprietary] = 3; + break; + case VK_VENDOR_ID_INTEL: + driver_priorities[vk::DriverId::eIntelOpenSourceMESA] = 1; + driver_priorities[vk::DriverId::eIntelProprietaryWindows] = 2; + break; + case VK_VENDOR_ID_NVIDIA: + driver_priorities[vk::DriverId::eNvidiaProprietary] = 1; +#if defined(VK_API_VERSION_1_3) && VK_HEADER_VERSION >= 235 + driver_priorities[vk::DriverId::eMesaNvk] = 2; +#endif + break; + } + + if (driver_priorities.count(old_driver.driverID)) { + old_priority = driver_priorities[old_driver.driverID]; + } + if (driver_priorities.count(new_driver.driverID)) { + new_priority = driver_priorities[new_driver.driverID]; + } + + if (new_priority < old_priority) { + auto r = std::remove(vk_instance.device_indices.begin(), vk_instance.device_indices.end(), *old_device); + vk_instance.device_indices.erase(r, vk_instance.device_indices.end()); + vk_instance.device_indices.push_back(i); + + VK_LOG_DEBUG("Prioritize device " << i << " driver " << new_driver.driverName << " over device " << *old_device << " driver " << old_driver.driverName); + } + else { + VK_LOG_DEBUG("Prioritize device " << *old_device << " driver " << old_driver.driverName << " over device " << i << " driver " << new_driver.driverName << std::endl); + } + } + } + } + + // If no dedicated GPUs found, fall back to GPU 0 + if (vk_instance.device_indices.empty()) { + vk_instance.device_indices.push_back(0); + } + } + GGML_LOG_DEBUG("ggml_vulkan: Found %zu Vulkan devices:\n", vk_instance.device_indices.size()); + + for (size_t i = 0; i < vk_instance.device_indices.size(); i++) { + ggml_vk_print_gpu_info(i); + } +} + +static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) { + VK_LOG_DEBUG("ggml_vk_init(" << ctx->name << ", " << idx << ")"); + ggml_vk_instance_init(); + GGML_ASSERT(idx < vk_instance.device_indices.size()); + + ctx->name = GGML_VK_NAME + std::to_string(idx); + + ctx->device = ggml_vk_get_device(idx); + + ctx->semaphore_idx = 0; + ctx->event_idx = 0; + + ctx->prealloc_size_x = 0; + ctx->prealloc_size_y = 0; + ctx->prealloc_size_split_k = 0; + + ctx->fence = ctx->device->device.createFence({}); + +#ifdef GGML_VULKAN_CHECK_RESULTS + const char* skip_checks = getenv("GGML_VULKAN_SKIP_CHECKS"); + vk_skip_checks = (skip_checks == NULL ? 0 : atoi(skip_checks)); + const char* output_tensor = getenv("GGML_VULKAN_OUTPUT_TENSOR"); + vk_output_tensor = (output_tensor == NULL ? 0 : atoi(output_tensor)); +#endif +} + +static vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type type) { + VK_LOG_DEBUG("ggml_vk_get_to_fp16()"); + switch (type) { + case GGML_TYPE_F32: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return nullptr; + } + + return ctx->device->pipeline_dequant[type]; +} + +static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec) { + VK_LOG_DEBUG("ggml_vk_get_mul_mat_mat_pipeline(" << ggml_type_name(src0_type) << ", " << ggml_type_name(src1_type) << ")"); + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_f32; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_f32_f16; + } + if (prec == GGML_PREC_DEFAULT && ctx->device->fp16 && !(ctx->device->coopmat_support && !ctx->device->coopmat_acc_f16_support)) { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_f16_f32.f16acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_f16.f16acc; + } + } else { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_f16_f32.f32acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_f16.f32acc; + } + } + + if (src1_type != GGML_TYPE_F32 && !ctx->device->coopmat2) { + return nullptr; + } + + switch (src0_type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return nullptr; + } + + if (ctx->device->coopmat2) { + assert(src1_type == GGML_TYPE_F16); + return ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f16acc; + } + return ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f32acc; +} + +static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * ctx, ggml_type a_type, ggml_type b_type, uint32_t num_cols) { + VK_LOG_DEBUG("ggml_vk_get_dequantize_mul_mat_vec()"); + GGML_ASSERT(b_type == GGML_TYPE_F32 || b_type == GGML_TYPE_F16); + GGML_ASSERT(num_cols >= 1 && num_cols <= mul_mat_vec_max_cols); + + switch (a_type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return nullptr; + } + + return b_type == GGML_TYPE_F32 ? ctx->device->pipeline_dequant_mul_mat_vec_f32_f32[a_type][num_cols-1] : ctx->device->pipeline_dequant_mul_mat_vec_f16_f32[a_type][num_cols-1]; +} + +static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec) { + VK_LOG_DEBUG("ggml_vk_get_mul_mat_mat_id_pipeline()"); + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_id_f32; + } + if (prec == GGML_PREC_DEFAULT && ctx->device->fp16 && !(ctx->device->coopmat_support && !ctx->device->coopmat_acc_f16_support)) { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_id_f16_f32.f16acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_id_f16.f16acc; + } + } else { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_id_f16_f32.f32acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_id_f16.f32acc; + } + } + + GGML_ASSERT(src1_type == GGML_TYPE_F32); + + switch (src0_type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return nullptr; + } + + return ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat_id[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat_id[src0_type].f32acc; +} + +static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context * ctx, ggml_type a_type, ggml_type b_type) { + VK_LOG_DEBUG("ggml_vk_get_dequantize_mul_mat_vec()"); + GGML_ASSERT(b_type == GGML_TYPE_F32); + + switch (a_type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return nullptr; + } + + return ctx->device->pipeline_dequant_mul_mat_vec_id_f32[a_type]; +} + +static vk_buffer ggml_vk_pool_malloc(ggml_backend_vk_context * ctx, size_t size) { + VK_LOG_DEBUG("ggml_vk_pool_malloc(" << size << ")"); + VK_LOG_MEMORY("ggml_vk_pool_malloc"); + + int best_i = -1; + size_t best_size = std::numeric_limits::max(); //smallest unused buffer that fits our needs + int worst_i = -1; + size_t worst_size = 0; //largest unused buffer seen so far + for (int i = 0; i < MAX_VK_BUFFERS; ++i) { + vk_buffer &b = ctx->buffer_pool[i]; + if (b != nullptr && b->size >= size && b->size < best_size) { + best_i = i; + best_size = b->size; + } + if (b != nullptr && b->size > worst_size) { + worst_i = i; + worst_size = b->size; + } + } + if(best_i != -1) { + //found the smallest buffer that fits our needs + vk_buffer b = ctx->buffer_pool[best_i]; + ctx->buffer_pool[best_i].reset(); + return b; + } + if(worst_i != -1) { + //no buffer that fits our needs, resize largest one to save memory + vk_buffer& b = ctx->buffer_pool[worst_i]; + ggml_vk_destroy_buffer(b); + } + + return ggml_vk_create_buffer_device(ctx->device, size); +} + +static void ggml_vk_pool_free(ggml_backend_vk_context * ctx, vk_buffer& buffer) { + VK_LOG_DEBUG("ggml_vk_pool_free(" << buffer->size << ")"); + for (int i = 0; i < MAX_VK_BUFFERS; ++i) { + vk_buffer& b = ctx->buffer_pool[i]; + if (b == nullptr) { + b = buffer; + return; + } + } + std::cerr << "ggml_vulkan: WARNING: vk buffer pool full, increase MAX_VK_BUFFERS" << std::endl; + ggml_vk_destroy_buffer(buffer); +} + +// Returns an available temporary buffer that may only be used temporarily, it will be reused +static vk_buffer ggml_vk_create_buffer_temp(ggml_backend_vk_context * ctx, size_t size) { + // Try to find existing temp buffer with enough capacity + for (auto& buffer : ctx->gc.temp_buffers) { + if (buffer->size >= size) { + return buffer; + } + } + + VK_LOG_MEMORY("ggml_vk_create_buffer_temp(" << size << ")"); + + // Otherwise create new buffer + vk_buffer buf = ggml_vk_pool_malloc(ctx, size); + ctx->gc.temp_buffers.push_back(buf); + + return buf; +} + +static void * ggml_vk_host_malloc(vk_device& device, size_t size) { + VK_LOG_MEMORY("ggml_vk_host_malloc(" << size << ")"); + vk_buffer buf = ggml_vk_create_buffer(device, size, + vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached, + vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); + + if(!(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible)) { + fprintf(stderr, "WARNING: failed to allocate %.2f MB of pinned memory\n", + size/1024.0/1024.0); + device->device.freeMemory(buf->device_memory); + device->device.destroyBuffer(buf->buffer); + return nullptr; + } + + device->pinned_memory.push_back(std::make_tuple(buf->ptr, size, buf)); + + return buf->ptr; +} + +static void ggml_vk_host_free(vk_device& device, void* ptr) { + if (ptr == nullptr) { + return; + } + VK_LOG_MEMORY("ggml_vk_host_free(" << ptr << ")"); + vk_buffer buf; + size_t index; + for (size_t i = 0; i < device->pinned_memory.size(); i++) { + const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]); + const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]); + if (ptr >= addr && ptr < endr) { + buf = std::get<2>(device->pinned_memory[i]); + index = i; + break; + } + } + if (buf == nullptr) { + fprintf(stderr, "WARNING: failed to free pinned memory: memory not in map\n"); + return; + } + + ggml_vk_destroy_buffer(buf); + + device->pinned_memory.erase(device->pinned_memory.begin() + index); +} + +static void ggml_vk_host_get(vk_device& device, const void * ptr, vk_buffer& buf, size_t& buf_offset) { + buf = nullptr; + buf_offset = 0; + for (size_t i = 0; i < device->pinned_memory.size(); i++) { + const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]); + const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]); + if (ptr >= addr && ptr < endr) { + buf = std::get<2>(device->pinned_memory[i]); + buf_offset = ((const uint8_t *)ptr) - addr; + break; + } + } +} + +static vk_submission ggml_vk_begin_submission(vk_device& device, vk_queue& q, bool one_time = true) { + vk_submission s; + s.buffer = ggml_vk_create_cmd_buffer(device, q); + if (one_time) { + s.buffer.begin({ vk::CommandBufferUsageFlagBits::eOneTimeSubmit }); + } else { + s.buffer.begin({ vk::CommandBufferUsageFlags{} }); + } + + return s; +} + + + +static void ggml_vk_dispatch_pipeline(ggml_backend_vk_context* ctx, vk_context& subctx, vk_pipeline& pipeline, std::initializer_list const& descriptor_buffer_infos, size_t push_constant_size, const void* push_constants, std::array elements) { + const uint32_t wg0 = CEIL_DIV(elements[0], pipeline->wg_denoms[0]); + const uint32_t wg1 = CEIL_DIV(elements[1], pipeline->wg_denoms[1]); + const uint32_t wg2 = CEIL_DIV(elements[2], pipeline->wg_denoms[2]); + VK_LOG_DEBUG("ggml_vk_dispatch_pipeline(" << pipeline->name << ", {"; + for (auto& buffer : descriptor_buffer_infos) { + std::cerr << "(" << buffer.buffer << ", " << buffer.offset << ", " << buffer.range << "), "; + } + std::cerr << "}, (" << wg0 << "," << wg1 << "," << wg2 << "))"); + GGML_ASSERT(pipeline->descriptor_set_idx < pipeline->descriptor_sets.size()); + GGML_ASSERT(descriptor_buffer_infos.size() == pipeline->parameter_count); + + vk::DescriptorSet& descriptor_set = pipeline->descriptor_sets[pipeline->descriptor_set_idx++]; + vk::WriteDescriptorSet write_descriptor_set{ descriptor_set, 0, 0, pipeline->parameter_count, vk::DescriptorType::eStorageBuffer, nullptr, descriptor_buffer_infos.begin() }; + ctx->device->device.updateDescriptorSets({ write_descriptor_set }, {}); + + subctx->s->buffer.pushConstants(pipeline->layout, vk::ShaderStageFlagBits::eCompute, 0, push_constant_size, push_constants); + subctx->s->buffer.bindPipeline(vk::PipelineBindPoint::eCompute, pipeline->pipeline); + subctx->s->buffer.bindDescriptorSets(vk::PipelineBindPoint::eCompute, + pipeline->layout, + 0, + { descriptor_set }, + {}); + subctx->s->buffer.dispatch(wg0, wg1, wg2); +} + +static void ggml_vk_end_submission(vk_submission& s, std::vector wait_semaphores, std::vector signal_semaphores) { + s.buffer.end(); + + s.wait_semaphores = std::move(wait_semaphores); + s.signal_semaphores = std::move(signal_semaphores); +} + +static void ggml_vk_ctx_end(vk_context& ctx) { + VK_LOG_DEBUG("ggml_vk_ctx_end(" << ctx << ", " << ctx->seqs.size() << ")"); + if (ctx->s == nullptr) { + return; + } + + ctx->s->buffer.end(); + ctx->s = nullptr; +} + +static void ggml_vk_ctx_begin(vk_device& device, vk_context& subctx) { + VK_LOG_DEBUG("ggml_vk_ctx_begin(" << device->name << ")"); + if (subctx->s != nullptr) { + ggml_vk_ctx_end(subctx); + } + + subctx->seqs.push_back({ ggml_vk_begin_submission(device, *subctx->q) }); + subctx->s = subctx->seqs[subctx->seqs.size() - 1].data(); +} + +static size_t ggml_vk_align_size(size_t width, size_t align) { + VK_LOG_DEBUG("ggml_vk_align_size(" << width << ", " << align << ")"); + return CEIL_DIV(width, align) * align; +} + +static void deferred_memcpy(void * dst, const void * src, size_t size, std::vector* memcpys = nullptr) { + if (memcpys == nullptr) { + memcpy(dst, src, size); + } else { + memcpys->emplace_back(dst, src, size); + } +} + +static void ggml_vk_ensure_sync_staging_buffer(vk_device& device, size_t size) { + if (device->sync_staging == nullptr || device->sync_staging->size < size) { + VK_LOG_MEMORY("ggml_vk_ensure_sync_staging_buffer(" << size << ")"); + ggml_vk_destroy_buffer(device->sync_staging); + device->sync_staging = ggml_vk_create_buffer_check(device, size, + vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached, + vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); + } +} + +static void ggml_vk_buffer_write_nc_async(ggml_backend_vk_context * ctx, vk_context& subctx, vk_buffer& dst, size_t offset, const ggml_tensor * tensor, bool sync_staging = false) { + VK_LOG_DEBUG("ggml_vk_buffer_write_nc_async(" << tensor << ")"); + GGML_ASSERT(!ggml_is_contiguous(tensor)); + // Buffer is already mapped + if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { + std::cerr << "ggml_vulkan: buffer_write_nc_async dst buffer is host_visible. Use synchronous write." << std::endl; + GGML_ABORT("fatal error"); + } + // Check if src is pinned memory + vk_buffer buf = nullptr; + size_t buf_offset = 0; + ggml_vk_host_get(ctx->device, tensor->data, buf, buf_offset); + + const uint64_t ne0 = tensor->ne[0]; + const uint64_t ne1 = tensor->ne[1]; + const uint64_t ne2 = tensor->ne[2]; + const uint64_t ne3 = tensor->ne[3]; + const uint64_t nb0 = tensor->nb[0]; + const uint64_t nb1 = tensor->nb[1]; + const uint64_t nb2 = tensor->nb[2]; + const uint64_t nb3 = tensor->nb[3]; + const ggml_type type = tensor->type; + const uint64_t ts = ggml_type_size(type); + const uint64_t bs = ggml_blck_size(type); + + const uint64_t dstnb0 = ts; + const uint64_t dstnb1 = dstnb0*(ne0/bs); + const uint64_t dstnb2 = dstnb1*ne1; + const uint64_t dstnb3 = dstnb2*ne2; + + const uint64_t ne = ggml_nelements(tensor); + + if (buf != nullptr) { + // Memory is pinned, use as staging buffer + std::vector slices; + + for (uint64_t i3 = 0; i3 < ne3; i3++) { + for (uint64_t i2 = 0; i2 < ne2; i2++) { + // Find longest contiguous slice + if (ne1*nb1 == dstnb2) { + slices.push_back({ buf_offset + i3*nb3 + i2*nb2, offset + i3*dstnb3 + i2*dstnb2, dstnb2 }); + } else { + for (uint64_t i1 = 0; i1 < ne1; i1++) { + if (ne0*nb0/bs == dstnb1) { + slices.push_back({ buf_offset + i3*nb3 + i2*nb2 + i1*nb1, offset + i3*dstnb3 + i2*dstnb2 + i1*dstnb1, dstnb1 }); + } else { + const uint64_t s_off = buf_offset + i3*nb3 + i2*nb2 + i1*nb1; + const uint64_t d_off = offset + i3*dstnb3 + i2*dstnb2 + i1*dstnb1; + for (uint64_t i0 = 0; i0 < ne0; i0++) { + slices.push_back({ s_off + i1*nb0, d_off + i0*dstnb0, dstnb0 }); + } + } + } + } + } + } + + ggml_vk_sync_buffers(subctx); + subctx->s->buffer.copyBuffer(buf->buffer, dst->buffer, slices); + return; + } + + if (!sync_staging) { + GGML_ABORT("Asynchronous write to non-pinned memory not supported"); + } + + // Staging buffer required + vk_buffer& staging = ctx->device->sync_staging; + const uint64_t copy_size = ts*ne/bs; + ggml_vk_ensure_sync_staging_buffer(ctx->device, copy_size); + VkBufferCopy buf_copy{ 0, offset, copy_size }; + + ggml_vk_sync_buffers(subctx); + vkCmdCopyBuffer(subctx->s->buffer, (VkBuffer)staging->buffer, (VkBuffer)dst->buffer, 1, &buf_copy); + + for (uint64_t i3 = 0; i3 < ne3; i3++) { + for (uint64_t i2 = 0; i2 < ne2; i2++) { + // Find longest contiguous slice + if (ne1*nb1 == dstnb2) { + deferred_memcpy((uint8_t *)staging->ptr + i3*dstnb3 + i2*dstnb2, (const uint8_t *) tensor->data + buf_offset + i3*nb3 + i2*nb2, dstnb2, &subctx->in_memcpys); + } else { + for (uint64_t i1 = 0; i1 < ne1; i1++) { + if (ne0*nb0/bs == dstnb1) { + deferred_memcpy((uint8_t *)staging->ptr + i3*dstnb3 + i2*dstnb2 + i1*dstnb1, (const uint8_t *) tensor->data + buf_offset + i3*nb3 + i2*nb2 + i1*nb1, dstnb1, &subctx->in_memcpys); + } else { + const uint64_t s_off = buf_offset + i3*nb3 + i2*nb2 + i1*nb1; + const uint64_t d_off = i3*dstnb3 + i2*dstnb2 + i1*dstnb1; + for (uint64_t i0 = 0; i0 < ne0; i0++) { + deferred_memcpy((uint8_t *)staging->ptr + d_off + i0*dstnb0, (const uint8_t *) tensor->data + s_off + i0*nb0, dstnb0, &subctx->in_memcpys); + } + } + } + } + } + } +} + +static void ggml_vk_buffer_write_2d_async(vk_context subctx, vk_buffer& dst, size_t offset, const void * src, size_t spitch, size_t width, size_t height, bool sync_staging = false) { + VK_LOG_DEBUG("ggml_vk_buffer_write_2d_async(" << width << ", " << height << ")"); + // Buffer is already mapped + if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { + std::cerr << "ggml_vulkan: buffer_write_async dst buffer is host_visible. Use synchronous write." << std::endl; + GGML_ABORT("fatal error"); + } + // Check if src is pinned memory + vk_buffer buf = nullptr; + size_t buf_offset = 0; + ggml_vk_host_get(dst->device, src, buf, buf_offset); + + if (buf != nullptr) { + // Memory is pinned, use as staging buffer + std::vector slices(1); + if (width == spitch) { + // Only do single write if stride is equal + slices[0].srcOffset = buf_offset; + slices[0].dstOffset = offset; + slices[0].size = width * height; + } else { + slices.resize(height); + for (size_t i = 0; i < height; i++) { + slices[i].srcOffset = buf_offset + i * spitch; + slices[i].dstOffset = offset + i * width; + slices[i].size = width; + } + } + + ggml_vk_sync_buffers(subctx); + subctx->s->buffer.copyBuffer(buf->buffer, dst->buffer, slices); + return; + } + VK_LOG_DEBUG("STAGING"); + + if (!sync_staging) { + GGML_ABORT("Asynchronous write to non-pinned memory not supported"); + } + + // Staging buffer required + const size_t copy_size = width*height; + ggml_vk_ensure_sync_staging_buffer(dst->device, copy_size); + + vk_buffer& staging_buffer = dst->device->sync_staging; + + VkBufferCopy buf_copy = { + 0, + offset, + copy_size}; + + ggml_vk_sync_buffers(subctx); + vkCmdCopyBuffer(subctx->s->buffer, (VkBuffer)staging_buffer->buffer, (VkBuffer)dst->buffer, 1, &buf_copy); + + if (width == spitch) { + deferred_memcpy((uint8_t *)staging_buffer->ptr, src, width * height, &subctx->in_memcpys); + } else { + for (size_t i = 0; i < height; i++) { + deferred_memcpy((uint8_t *)staging_buffer->ptr + i * width, (const uint8_t *) src + i * spitch, width, &subctx->in_memcpys); + } + } +} + +static void ggml_vk_buffer_write_async(vk_context subctx, vk_buffer& dst, size_t offset, const void * src, size_t size, bool sync_staging = false) { + VK_LOG_DEBUG("ggml_vk_buffer_write_async(" << size << ")"); + return ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, size, size, 1, sync_staging); +} + +static void ggml_vk_buffer_write_2d(vk_buffer& dst, size_t offset, const void * src, size_t spitch, size_t width, size_t height) { + VK_LOG_DEBUG("ggml_vk_buffer_write_2d(" << width << ", " << height << ")"); + // Buffer is already mapped + if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { + GGML_ASSERT(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent); + + for (size_t i = 0; i < height; i++) { + memcpy((uint8_t *)dst->ptr + offset + i * width, (const uint8_t *) src + i * spitch, width); + } + } else { + vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue); + ggml_vk_ctx_begin(dst->device, subctx); + ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, spitch, width, height, true); + ggml_vk_ctx_end(subctx); + + for (auto& cpy : subctx->in_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + + ggml_vk_submit(subctx, dst->device->fence); + VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_buffer_write_2d waitForFences"); + dst->device->device.resetFences({ dst->device->fence }); + } +} + +static void ggml_vk_buffer_write(vk_buffer& dst, size_t offset, const void * src, size_t size) { + VK_LOG_DEBUG("ggml_vk_buffer_write(" << size << ")"); + ggml_vk_buffer_write_2d(dst, offset, src, 0, size, 1); +} + +static void ggml_vk_buffer_read_2d_async(vk_context subctx, vk_buffer& src, size_t offset, void * dst, size_t spitch, size_t dpitch, size_t width, size_t height, bool sync_staging = false) { + VK_LOG_DEBUG("ggml_vk_buffer_read_2d_async(offset=" << offset << ", width=" << width << ", height=" << height << ")"); + GGML_ASSERT(width > 0); + GGML_ASSERT(height > 0); + GGML_ASSERT(src != nullptr); + + // TODO: staging_offset is not used + + // Check if dst is pinned memory + vk_buffer buf = nullptr; + size_t buf_offset = 0; + ggml_vk_host_get(src->device, dst, buf, buf_offset); + + std::vector slices(1); + if (width == spitch && width == dpitch) { + // Only do single write if stride is equal + slices[0].srcOffset = offset; + slices[0].dstOffset = buf_offset; + slices[0].size = width * height; + } else { + slices.resize(height); + for (size_t i = 0; i < height; i++) { + slices[i].srcOffset = offset + i * spitch; + slices[i].dstOffset = buf_offset + i * dpitch; + slices[i].size = width; + } + } + + if (buf != nullptr) { + // Memory is pinned, use as staging buffer + ggml_vk_sync_buffers(subctx); + subctx->s->buffer.copyBuffer(src->buffer, buf->buffer, slices); + + return; + } + VK_LOG_DEBUG("STAGING"); + + if (!sync_staging) { + GGML_ABORT("Asynchronous read from non-pinned memory not supported"); + } + + // Fall back to staging buffer + const size_t copy_size = dpitch * height; + ggml_vk_ensure_sync_staging_buffer(src->device, copy_size); + + vk_buffer& staging_buffer = src->device->sync_staging; + + ggml_vk_sync_buffers(subctx); + subctx->s->buffer.copyBuffer(src->buffer, staging_buffer->buffer, slices); + + deferred_memcpy(dst, staging_buffer->ptr, copy_size, &subctx->out_memcpys); +} + +static void ggml_vk_buffer_read_async(vk_context subctx, vk_buffer& src, size_t offset, void * dst, size_t size, bool sync_staging = false) { + return ggml_vk_buffer_read_2d_async(subctx, src, offset, dst, size, size, size, 1, sync_staging); +} + +static void ggml_vk_buffer_read(vk_buffer& src, size_t offset, void * dst, size_t size) { + VK_LOG_DEBUG("ggml_vk_buffer_read(" << src->buffer << ", " << offset << ", " << size << ")"); + + // If the device is not an UMA device the memory is host-accessible through rebar. While writing + // through PCIe is sufficient fast reading back data from PCIe is slower than going through + // the HW device to host copy path. + if(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible && src->device->uma) { + GGML_ASSERT(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent); + + memcpy(dst, (uint8_t *) src->ptr + offset, size); + } else { + vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue); + ggml_vk_ctx_begin(src->device, subctx); + ggml_vk_buffer_read_async(subctx, src, offset, dst, size, true); + ggml_vk_ctx_end(subctx); + + ggml_vk_submit(subctx, src->device->fence); + VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_read waitForFences"); + src->device->device.resetFences({ src->device->fence }); + + for (auto& cpy : subctx->out_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + } +} + +static void ggml_vk_buffer_copy_async(vk_context& ctx, vk_buffer& dst, size_t dst_offset, vk_buffer& src, size_t src_offset, size_t size) { + VK_LOG_DEBUG("ggml_vk_buffer_copy_async(" << size << ")"); + // Make sure both buffers are on same device + GGML_ASSERT(src->device == dst->device); + + VkBufferCopy bc{ src_offset, dst_offset, size }; + + vkCmdCopyBuffer(ctx->s->buffer, (VkBuffer)src->buffer, (VkBuffer)dst->buffer, 1, &bc); +} + +static void ggml_vk_buffer_copy(vk_buffer& dst, size_t dst_offset, vk_buffer& src, size_t src_offset, size_t size) { + if (src->device == dst->device) { + VK_LOG_DEBUG("ggml_vk_buffer_copy(SINGLE_DEVICE, " << size << ")"); + // Copy within the device + vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue); + ggml_vk_ctx_begin(src->device, subctx); + ggml_vk_buffer_copy_async(subctx, dst, dst_offset, src, src_offset, size); + ggml_vk_ctx_end(subctx); + ggml_vk_submit(subctx, src->device->fence); + VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_copy waitForFences"); + src->device->device.resetFences({ src->device->fence }); + } else { + VK_LOG_DEBUG("ggml_vk_buffer_copy(MULTI_DEVICE, " << size << ")"); + // Copy device to device + ggml_vk_ensure_sync_staging_buffer(src->device, size); + ggml_vk_ensure_sync_staging_buffer(dst->device, size); + + // Copy to src staging buffer + ggml_vk_buffer_copy(src->device->sync_staging, 0, src, src_offset, size); + // memcpy to dst staging buffer + memcpy(dst->device->sync_staging->ptr, src->device->sync_staging->ptr, size); + // Copy to dst buffer + ggml_vk_buffer_copy(dst, dst_offset, dst->device->sync_staging, 0, size); + } +} + +static void ggml_vk_buffer_memset(vk_buffer& dst, size_t offset, uint32_t c, size_t size) { + VK_LOG_DEBUG("ggml_vk_buffer_memset(" << offset << ", " << c << ", " << size << ")"); + + vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue); + ggml_vk_ctx_begin(dst->device, subctx); + subctx->s->buffer.fillBuffer(dst->buffer, offset, size, c); + ggml_vk_ctx_end(subctx); + + ggml_vk_submit(subctx, dst->device->fence); + VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_memset waitForFences"); + dst->device->device.resetFences({ dst->device->fence }); +} + +static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, int m, int n, int k, const vk_pipeline& pipeline) { + VK_LOG_DEBUG("ggml_vk_guess_split_k(" << m << ", " << n << ", " << k << ")"); + + uint32_t split_k = 1; + if (ctx->device->shader_core_count != 0 && m >= (int)pipeline->wg_denoms[0] && n >= (int)pipeline->wg_denoms[1]) { + // If k is 'large' and the SMs will fill less than halfway, use split_k. + uint32_t m_tiles = CEIL_DIV(m, pipeline->wg_denoms[0]); + uint32_t n_tiles = CEIL_DIV(n, pipeline->wg_denoms[1]); + if (k >= 2048 && m_tiles * n_tiles < ctx->device->shader_core_count / 2) { + split_k = ctx->device->shader_core_count / (m_tiles * n_tiles); + // Clamp to 2 or 4 + split_k = std::min(split_k, 4u); + if (split_k == 3) { + split_k = 2; + } + } + } + + return split_k; +} + +static vk_pipeline ggml_vk_guess_matmul_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, bool aligned) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline(" << m << ", " << n << ", " << aligned << ")"); + + if (ctx->device->coopmat2) { + if ((ctx->device->mul_mat_l && (m % mmp->l->wg_denoms[0]) == 0 && (n % mmp->l->wg_denoms[1]) == 0) || (!ctx->device->mul_mat_m && !ctx->device->mul_mat_s)) { + return aligned ? mmp->a_l : mmp->l; + } + if ((ctx->device->mul_mat_m && (m % mmp->m->wg_denoms[0]) == 0 && (n % mmp->m->wg_denoms[1]) == 0) || !ctx->device->mul_mat_s) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_s : mmp->s; + } + + if ((ctx->device->mul_mat_s && (m <= 32 || n <= 32)) || (!ctx->device->mul_mat_m && !ctx->device->mul_mat_l)) { + return aligned ? mmp->a_s : mmp->s; + } + if ((ctx->device->mul_mat_m && (m <= 64 || n <= 64)) || !ctx->device->mul_mat_l) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_l : mmp->l; +} + +static uint32_t ggml_vk_guess_matmul_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ")"); + return ggml_vk_guess_matmul_pipeline(ctx, mmp, m, n, true)->align; +} + +static void ggml_vk_matmul( + ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, + vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& split_k_buffer, + uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, + uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, + uint32_t split_k, uint32_t batch, uint32_t ne02, uint32_t ne12, uint32_t broadcast2, uint32_t broadcast3) { + VK_LOG_DEBUG("ggml_vk_matmul(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), split_k: (" << (split_k_buffer.buffer != nullptr ? split_k_buffer.buffer->buffer : VK_NULL_HANDLE) << ", " << split_k_buffer.offset << ", " << split_k_buffer.size << "), m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", split_k: " << split_k << ", batch: " << batch << ", ne02: " << ne02 << ", ne12: " << ne12 << ", broadcast2: " << broadcast2 << ", broadcast3: " << broadcast3 << ")"); + ggml_vk_sync_buffers(subctx); + if (split_k == 1) { + const vk_mat_mat_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, k, ne02, ne12, broadcast2, broadcast3 }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d }, sizeof(vk_mat_mat_push_constants), &pc, { m, n, batch }); + return; + } + + GGML_ASSERT(batch_stride_d == m * n); + + const vk_mat_mat_push_constants pc1 = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, CEIL_DIV(k, split_k), ne02, ne12, broadcast2, broadcast3 }; + // Make sure enough workgroups get assigned for split k to work + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, split_k_buffer }, sizeof(vk_mat_mat_push_constants), &pc1, { (CEIL_DIV(m, pipeline->wg_denoms[0]) * pipeline->wg_denoms[0]) * split_k, n, batch }); + ggml_vk_sync_buffers(subctx); + const std::array pc2 = { (uint32_t)(m * n * batch), split_k }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_matmul_split_k_reduce, { split_k_buffer, d }, pc2.size() * sizeof(uint32_t), pc2.data(), { m * n * batch, 1, 1 }); +} + +static vk_pipeline ggml_vk_guess_matmul_id_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, bool aligned) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline(" << m << ", " << n << ", " << aligned << ")"); + + if (ctx->device->coopmat2) { + if ((ctx->device->mul_mat_id_l && (m % mmp->l->wg_denoms[0]) == 0 && (n % mmp->l->wg_denoms[1]) == 0) || (!ctx->device->mul_mat_id_m && !ctx->device->mul_mat_id_s)) { + return aligned ? mmp->a_l : mmp->l; + } + if ((ctx->device->mul_mat_id_m && (m % mmp->m->wg_denoms[0]) == 0 && (n % mmp->m->wg_denoms[1]) == 0) || !ctx->device->mul_mat_id_s) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_s : mmp->s; + } + + if ((ctx->device->mul_mat_id_s && (m <= 32 || n <= 32)) || (!ctx->device->mul_mat_id_m && !ctx->device->mul_mat_id_l)) { + return aligned ? mmp->a_s : mmp->s; + } + if ((ctx->device->mul_mat_id_m && (m <= 64 || n <= 64)) || !ctx->device->mul_mat_id_l) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_l : mmp->l; +} + +static uint32_t ggml_vk_guess_matmul_id_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ")"); + return ggml_vk_guess_matmul_id_pipeline(ctx, mmp, m, n, true)->align; +} + +static void ggml_vk_matmul_id( + ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, + vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& ids, + uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, + uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, + uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11) { + VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), " << + "m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " << + "batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " << + "n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")"); + ggml_vk_sync_buffers(subctx); + const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, + nei0, nei1, nbi1, ne11 }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids }, sizeof(vk_mat_mat_id_push_constants), &pc, { m, nei1, n_as }); +} + +static bool ggml_vk_dim01_contiguous(const ggml_tensor * tensor) { + return + tensor->nb[0] == ggml_type_size(tensor->type) && + tensor->nb[1] == (tensor->nb[0]*tensor->ne[0])/ggml_blck_size(tensor->type) && + tensor->nb[3] == tensor->nb[2]*tensor->ne[2]; +} + +static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src, const ggml_tensor * dst, ggml_type to) { + + // Choose "contiguous copy" shader if src/dst are contiguous + bool contig = ggml_is_contiguous(src) && (!dst || ggml_is_contiguous(dst)); + + if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F32) { + if (contig) { + return ctx->device->pipeline_contig_cpy_f32_f32; + } else { + return ctx->device->pipeline_cpy_f32_f32; + } + } + if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F16) { + if (contig) { + return ctx->device->pipeline_contig_cpy_f32_f16; + } else { + return ctx->device->pipeline_cpy_f32_f16; + } + } + if (src->type == GGML_TYPE_F16 && to == GGML_TYPE_F16) { + if (contig) { + return ctx->device->pipeline_contig_cpy_f16_f16; + } else { + return ctx->device->pipeline_cpy_f16_f16; + } + } + + std::cerr << "Missing CPY op for types: " << ggml_type_name(src->type) << " " << ggml_type_name(to) << std::endl; + GGML_ABORT("fatal error"); +} + +static void ggml_vk_cpy_to_contiguous(ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline pipeline, const ggml_tensor * tensor, vk_subbuffer&& in, vk_subbuffer&& out) { + VK_LOG_DEBUG("ggml_vk_cpy_to_contiguous((" << tensor << ", type=" << tensor->type << ", ne0=" << tensor->ne[0] << ", ne1=" << tensor->ne[1] << ", ne2=" << tensor->ne[2] << ", ne3=" << tensor->ne[3] << ", nb0=" << tensor->nb[0] << ", nb1=" << tensor->nb[1] << ", nb2=" << tensor->nb[2] << ", nb3=" << tensor->nb[3] << "), "; + std::cerr << "buffer in size=" << in.buffer->size << ", buffer out size=" << out.buffer->size << ")"); + const int tensor_type_size = ggml_type_size(tensor->type); + + const uint32_t ne = ggml_nelements(tensor); + std::array elements; + + if (ne > 262144) { + elements = { 512, 512, CEIL_DIV(ne, 262144) }; + } else if (ne > 512) { + elements = { 512, CEIL_DIV(ne, 512), 1 }; + } else { + elements = { ne, 1, 1 }; + } + + vk_op_unary_push_constants pc = { + (uint32_t)ne, + (uint32_t)tensor->ne[0], (uint32_t)tensor->ne[1], (uint32_t)tensor->ne[2], (uint32_t)tensor->ne[3], (uint32_t)tensor->nb[0] / tensor_type_size, (uint32_t)tensor->nb[1] / tensor_type_size, (uint32_t)tensor->nb[2] / tensor_type_size, (uint32_t)tensor->nb[3] / tensor_type_size, + (uint32_t)tensor->ne[0], (uint32_t)tensor->ne[1], (uint32_t)tensor->ne[2], (uint32_t)tensor->ne[3], 1 , (uint32_t)tensor->ne[0] , (uint32_t)(tensor->ne[0] * tensor->ne[1]) , (uint32_t)(tensor->ne[0] * tensor->ne[1] * tensor->ne[2]), + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }; + init_pushconst_fastdiv(pc); + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { in, out }, sizeof(vk_op_unary_push_constants), &pc, elements); +} + +static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); + GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT + GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + const uint64_t ne03 = src0->ne[3]; + + const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + const uint64_t ne13 = src1->ne[3]; + + const uint64_t ne20 = dst->ne[0]; + const uint64_t ne21 = dst->ne[1]; + + const uint64_t r2 = ne12 / ne02; + const uint64_t r3 = ne13 / ne03; + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + + vk_buffer d_Qx = nullptr; + size_t qx_buf_offset = 0; + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + + bool src0_uma = false; + bool src1_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + src0_uma = d_Qx != nullptr; + src1_uma = d_Qy != nullptr; + } + + const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); + // Reformat and convert to fp16 if src1 is non-contiguous, or for coopmat2 for better perf + const bool y_non_contig = (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || + !ggml_vk_dim01_contiguous(src1); + + const bool y_f32_kernel = src1->type == GGML_TYPE_F32 && !y_non_contig; + + vk_matmul_pipeline mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, src0->type, y_non_contig ? GGML_TYPE_F16 : src1->type, (ggml_prec)dst->op_params[0]); + + const bool qx_needs_dequant = mmp == nullptr || x_non_contig; + const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !y_f32_kernel) || y_non_contig; + + if (qx_needs_dequant) { + // Fall back to dequant + f16 mulmat + mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, GGML_TYPE_F16, y_f32_kernel ? GGML_TYPE_F32 : GGML_TYPE_F16, (ggml_prec)dst->op_params[0]); + } + + // Not implemented + GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT + + const int x_ne = ne01 * ne00; + const int y_ne = ne11 * ne10; + const int d_ne = ne11 * ne01; + + const uint32_t kpad = ggml_vk_align_size(ne10, ggml_vk_guess_matmul_pipeline_align(ctx, mmp, ne01, ne11)); + const bool aligned = ne10 == kpad && ne01 > 8 && ne11 > 8; + + vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline(ctx, mmp, ne01, ne11, aligned); + + const uint32_t split_k = ggml_vk_guess_split_k(ctx, ne01, ne11, ne10, pipeline); + + const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type); + const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); + const uint64_t x_sz = !qx_needs_dequant ? qx_sz : sizeof(ggml_fp16_t) * x_ne; + const uint64_t y_sz = y_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; + const uint64_t d_sz = sizeof(float) * d_ne; + + vk_pipeline to_fp16_vk_0 = nullptr; + vk_pipeline to_fp16_vk_1 = nullptr; + + if (x_non_contig) { + to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, GGML_TYPE_F16); + } else { + to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, src0->type); + } + if (y_non_contig) { + to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, GGML_TYPE_F16); + } else { + to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); + } + GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT + GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT + + if (dryrun) { + const uint64_t x_sz_upd = x_sz * ne02 * ne03; + const uint64_t y_sz_upd = y_sz * ne12 * ne13; + const uint64_t split_k_size = split_k > 1 ? d_sz * ne12 * ne13 * split_k : 0; + if ( + (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || + (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size) || + (split_k > 1 && split_k_size > ctx->device->max_memory_allocation_size)) { + GGML_ABORT("Requested preallocation size is too large"); + } + if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { + ctx->prealloc_size_x = x_sz_upd; + } + if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { + ctx->prealloc_size_y = y_sz_upd; + } + if (split_k > 1 && ctx->prealloc_size_split_k < split_k_size) { + ctx->prealloc_size_split_k = split_k_size; + } + + // Request descriptor sets + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + if (qx_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); + } + if (qy_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); + } + if (split_k > 1) { + ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, 1); + } + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + GGML_ASSERT(d_D->size >= d_buf_offset + d_sz * ne02 * ne03); + vk_buffer d_X; + uint64_t x_buf_offset = 0; + vk_buffer d_Y; + uint64_t y_buf_offset = 0; + if (!src0_uma) { + d_Qx = src0_buf_ctx->dev_buffer; + qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + if (!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qy != nullptr); + } + if (qx_needs_dequant) { + d_X = ctx->prealloc_x; + GGML_ASSERT(d_X->size >= x_sz * ne02 * ne03); + } else { + d_X = d_Qx; + x_buf_offset = qx_buf_offset; + GGML_ASSERT(qx_sz == x_sz); + } + if (qy_needs_dequant) { + d_Y = ctx->prealloc_y; + GGML_ASSERT(d_Y->size >= y_sz * ne02 * ne03); + } else { + d_Y = d_Qy; + y_buf_offset = qy_buf_offset; + GGML_ASSERT(qy_sz == y_sz); + } + + if (x_non_contig) { + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); + } else if (qx_needs_dequant) { + const std::vector pc = { (uint32_t)ne01, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)(ggml_nelements(src0)) }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz * ne02 * ne03 }, vk_subbuffer{ d_X, 0, x_sz * ne02 * ne03 } }, pc.size() * sizeof(uint32_t), pc.data(), { (uint32_t)(x_ne * ne02 * ne03), 1, 1}); + } + if (y_non_contig) { + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); + } + + uint32_t stride_batch_x = ne00*ne01; + uint32_t stride_batch_y = ne10*ne11; + + if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { + stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); + } + + if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { + stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); + } + + // compute + ggml_vk_matmul( + ctx, subctx, pipeline, + { d_X, x_buf_offset, x_sz * ne02 * ne03 }, { d_Y, y_buf_offset, y_sz * ne12 * ne13 }, + { d_D, d_buf_offset, d_sz * ne12 * ne13 }, { ctx->prealloc_split_k, 0, d_sz * ne12 * ne13 * split_k }, + ne01, ne11, ne10, + ne10, ne10, ne01, stride_batch_x, stride_batch_y, ne20*ne21, + split_k, ne12*ne13, ne02, ne12, r2, r3 + ); // NOLINT +} + +static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_vec_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << "),)"); + GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT + GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + const uint64_t ne03 = src0->ne[3]; + + const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + const uint64_t ne13 = src1->ne[3]; + + const uint64_t ne20 = dst->ne[0]; + const uint64_t ne21 = dst->ne[1]; + const uint64_t ne22 = dst->ne[2]; + const uint64_t ne23 = dst->ne[3]; + + const uint64_t r2 = ne12 / ne02; + const uint64_t r3 = ne13 / ne03; + + // batch_n indicates that we need to compute a few vector results, and this assumes + // ne12 and ne13 are 1. It overloads the batch_strides to hold the row strides. + GGML_ASSERT(ne11 == 1 || ne12 * ne13 == 1); + bool batch_n = ne11 > 1; + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + + vk_buffer d_Qx = nullptr; + size_t qx_buf_offset = 0; + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + + bool src0_uma = false; + bool src1_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + src0_uma = d_Qx != nullptr; + src1_uma = d_Qy != nullptr; + } + + const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); + const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); + + const bool f16_f32_kernel = src1->type == GGML_TYPE_F32; + + const bool qx_needs_dequant = x_non_contig; + const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !f16_f32_kernel) || y_non_contig; + + // Not implemented + GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT + + const uint64_t x_ne = ne01 * ne00; + const uint64_t y_ne = ne11 * ne10; + const uint64_t d_ne = ne11 * ne01; + + const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); + const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); + const uint64_t x_sz = x_non_contig ? ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment) : qx_sz; + const uint64_t y_sz = f16_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; + const uint64_t d_sz = sizeof(float) * d_ne; + + vk_pipeline to_fp16_vk_0 = nullptr; + vk_pipeline to_fp16_vk_1 = nullptr; + if (x_non_contig) { + to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, src0->type); + } + if (y_non_contig) { + to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, src1->type); + } else { + to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); + } + vk_pipeline dmmv = ggml_vk_get_dequantize_mul_mat_vec(ctx, src0->type, src1->type, ne11); + GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT + GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT + GGML_ASSERT(dmmv != nullptr); + + if (dryrun) { + const uint64_t x_sz_upd = x_sz * ne02 * ne03; + const uint64_t y_sz_upd = y_sz * ne12 * ne13; + if ( + (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || + (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { + GGML_ABORT("Requested preallocation size is too large"); + } + if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { + ctx->prealloc_size_x = x_sz_upd; + } + if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { + ctx->prealloc_size_y = y_sz_upd; + } + + // Request descriptor sets + if (qx_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); + } + if (qy_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); + } + ggml_pipeline_request_descriptor_sets(ctx->device, dmmv, 1); + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + vk_buffer d_X; + uint64_t x_buf_offset = 0; + vk_buffer d_Y; + uint64_t y_buf_offset = 0; + if(!src0_uma) { + d_Qx = src0_buf_ctx->dev_buffer; + qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + if(!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qy != nullptr); + } + if (qx_needs_dequant) { + d_X = ctx->prealloc_x; + } else { + d_X = d_Qx; + x_buf_offset = qx_buf_offset; + GGML_ASSERT(qx_sz == x_sz); + } + if (qy_needs_dequant) { + d_Y = ctx->prealloc_y; + } else { + d_Y = d_Qy; + y_buf_offset = qy_buf_offset; + GGML_ASSERT(qy_sz == y_sz); + } + + if (x_non_contig) { + GGML_ASSERT(x_sz == ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment)); + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); + } + if (y_non_contig) { + GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); + } + + // For batch_n, the A matrix is the same for each batch, and B/D use the row stride as the batch stride + uint32_t stride_batch_x = batch_n ? 0 : ne00*ne01; + uint32_t stride_batch_y = batch_n ? ne10 : (ne10*ne11); + uint32_t stride_batch_d = batch_n ? ne20 : (ne20*ne21); + + if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { + stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); + } + + if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { + stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); + } + + const uint32_t max_groups_x = ctx->device->properties.limits.maxComputeWorkGroupCount[0]; + + uint32_t groups_x = ne01; + uint32_t groups_z = 1; + + if (ne01 > max_groups_x) { + groups_z = 64; + groups_x = CEIL_DIV(groups_x, groups_z); + } + + // compute + const vk_mat_vec_push_constants pc = { + (uint32_t)ne00, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne01, + stride_batch_x, stride_batch_y, stride_batch_d, + (uint32_t)ne02, (uint32_t)ne12, (uint32_t)r2, (uint32_t)r3, + }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, dmmv, + { vk_subbuffer{ d_X, x_buf_offset, x_sz * ne02 * ne03 }, vk_subbuffer{ d_Y, y_buf_offset, y_sz * ne12 * ne13 }, vk_subbuffer{ d_D, d_buf_offset, d_sz * ne22 * ne23} }, + sizeof(vk_mat_vec_push_constants), &pc, { groups_x, (uint32_t)(ne12 * ne13), groups_z }); +} + +static void ggml_vk_mul_mat_vec_p021_f16_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_p021_f16_f32(" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); + GGML_ASSERT(ggml_is_permuted(src0) && ggml_is_permuted(src1)); + GGML_ASSERT(src0->nb[0] <= src0->nb[1] && src0->nb[2] <= src0->nb[3]); // NOLINT + GGML_ASSERT(src1->nb[0] <= src1->nb[1] && src1->nb[2] <= src1->nb[3]); // NOLINT + GGML_ASSERT(src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + // const uint64_t ne03 = src0->ne[3]; + + const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + // const uint64_t ne13 = src1->ne[3]; + + GGML_ASSERT(ne11 == 1); + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + + bool src1_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + src1_uma = d_Qy != nullptr; + } + + const uint64_t x_ne = ne00 * ne01 * ne02; + const uint64_t y_ne = ne10 * ne11 * ne12; + const uint64_t d_ne = ne01 * ne11 * ne12; + + const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); + const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); + const uint64_t d_sz = sizeof(float) * d_ne; + + if (dryrun) { + // Request descriptor sets + ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_mul_mat_vec_p021_f16_f32, 1); + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + vk_buffer d_Qx = src0_buf_ctx->dev_buffer; + const uint64_t qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + if (!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + + const uint64_t qy_buffer_offset = (qy_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; + const uint64_t qy_shader_offset = qy_buf_offset - qy_buffer_offset; + + const uint64_t d_buffer_offset = (d_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; + const uint64_t d_shader_offset = d_buf_offset - d_buffer_offset; + + // compute + const std::array pc = { (uint32_t)ne00, (uint32_t)ne01, (uint32_t)ne02, (uint32_t)ne12, (uint32_t)(qy_shader_offset / ggml_type_size(src1->type)), (uint32_t)(d_shader_offset / ggml_type_size(dst->type)) }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_mul_mat_vec_p021_f16_f32, { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz }, vk_subbuffer{ d_Qy, qy_buffer_offset, qy_sz + qy_shader_offset }, vk_subbuffer{ d_D, d_buffer_offset, d_sz + d_shader_offset } }, 6 * sizeof(uint32_t), &pc, { 1, (uint32_t)ne01, (uint32_t)ne12 }); +} + +static void ggml_vk_mul_mat_vec_nc_f16_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_nc_f16_f32((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); + GGML_ASSERT(!ggml_is_transposed(src0)); + GGML_ASSERT(!ggml_is_transposed(src1)); + GGML_ASSERT(!ggml_is_permuted(src0)); + GGML_ASSERT(src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + // const uint64_t ne03 = src0->ne[3]; + + const uint64_t nb01 = src0->nb[1]; + const uint64_t nb02 = src0->nb[2]; + + // const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + // const uint64_t ne13 = src1->ne[3]; + + GGML_ASSERT(ne11 == 1); + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + + bool src1_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + src1_uma = d_Qy != nullptr; + } + + const uint64_t d_ne = ne01 * ne11 * ne12; + + const uint32_t row_stride_x = nb01 / sizeof(ggml_fp16_t); + const uint32_t channel_stride_x = nb02 / sizeof(ggml_fp16_t); + + const uint64_t qx_sz = ggml_nbytes(src0); + const uint64_t qy_sz = ggml_nbytes(src1); + const uint64_t d_sz = sizeof(float) * d_ne; + + if (dryrun) { + // Request descriptor sets + ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_mul_mat_vec_nc_f16_f32, 1); + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + vk_buffer d_Qx = src0_buf_ctx->dev_buffer; + const uint64_t qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + if (!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + + const uint64_t qy_buffer_offset = (qy_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; + const uint64_t qy_shader_offset = qy_buf_offset - qy_buffer_offset; + + const uint64_t d_buffer_offset = (d_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; + const uint64_t d_shader_offset = d_buf_offset - d_buffer_offset; + + // compute + const std::array pc = { (uint32_t)ne00, (uint32_t)ne01, row_stride_x, channel_stride_x, (uint32_t)(ne12 / ne02), (uint32_t)(qy_shader_offset / ggml_type_size(src1->type)), (uint32_t)(d_shader_offset / ggml_type_size(dst->type)) }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_mul_mat_vec_nc_f16_f32, + { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz }, vk_subbuffer{ d_Qy, qy_buffer_offset, qy_sz + qy_shader_offset }, vk_subbuffer{ d_D, d_buffer_offset, d_sz + d_shader_offset } }, 7 * sizeof(uint32_t), &pc, { 1, (uint32_t)ne01, (uint32_t)ne12 }); +} + +static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat(" << src0 << ", " << src1 << ", " << dst << ")"); + if (src0->type == GGML_TYPE_F16 && ggml_is_permuted(src0) && ggml_is_permuted(src1) && dst->ne[1] == 1 && + // detect 0213 permutation, and batch size of 1 + src0->nb[0] <= src0->nb[2] && + src0->nb[2] <= src0->nb[1] && + src0->nb[1] <= src0->nb[3] && + src1->nb[0] <= src1->nb[2] && + src1->nb[2] <= src1->nb[1] && + src1->nb[1] <= src1->nb[3] && + src0->ne[3] == 1 && + src1->ne[3] == 1) { + ggml_vk_mul_mat_vec_p021_f16_f32(ctx, subctx, src0, src1, dst, dryrun); + } else if (src0->type == GGML_TYPE_F16 && !ggml_is_contiguous(src0) && !ggml_is_transposed(src1) && dst->ne[1] == 1 && + !ggml_is_permuted(src0) && !ggml_is_permuted(src1)) { + ggml_vk_mul_mat_vec_nc_f16_f32(ctx, subctx, src0, src1, dst, dryrun); + // mul_mat_vec supports batching ne12*ne13 when ne11==1, or treating ne11 as the batch size (up to four) + // when ne12 and ne13 are one. + } else if ((dst->ne[1] == 1 || (dst->ne[1] <= mul_mat_vec_max_cols && src1->ne[2] * src1->ne[3] == 1)) && + (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type))) { + ggml_vk_mul_mat_vec_q_f16(ctx, subctx, src0, src1, dst, dryrun); + } else { + ggml_vk_mul_mat_q_f16(ctx, subctx, src0, src1, dst, dryrun); + } +} + +static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_id_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << ids << ", name=" << ids->name << ", type=" << ids->type << ", ne0=" << ids->ne[0] << ", ne1=" << ids->ne[1] << ", ne2=" << ids->ne[2] << ", ne3=" << ids->ne[3] << ", nb0=" << ids->nb[0] << ", nb1=" << ids->nb[1] << ", nb2=" << ids->nb[2] << ", nb3=" << ids->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3] << "),)"); + GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT + GGML_ASSERT(ids->type == GGML_TYPE_I32); + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + const uint64_t ne03 = src0->ne[3]; + + const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + const uint64_t ne13 = src1->ne[3]; + + const uint64_t nei0 = ids->ne[0]; + const uint64_t nei1 = ids->ne[1]; + GGML_ASSERT(nei0 * nei1 <= 3072); + + const uint32_t nbi1 = ids->nb[1]; + const uint32_t nbi2 = ids->nb[2]; + + const uint64_t ne20 = dst->ne[0]; + const uint64_t ne21 = dst->ne[1]; + const uint64_t ne22 = dst->ne[2]; + const uint64_t ne23 = dst->ne[3]; + + const uint64_t n_as = ne02; + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + ggml_backend_vk_buffer_context * ids_buf_ctx = (ggml_backend_vk_buffer_context *)ids->buffer->context; + + vk_buffer d_Qx = nullptr; + size_t qx_buf_offset = 0; + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + vk_buffer d_ids = nullptr; + size_t ids_buf_offset = 0; + + bool src0_uma = false; + bool src1_uma = false; + bool ids_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + ggml_vk_host_get(ctx->device, ids->data, d_ids, ids_buf_offset); + src0_uma = d_Qx != nullptr; + src1_uma = d_Qy != nullptr; + ids_uma = d_ids != nullptr; + } + + const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); + const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); + + const bool y_f32_kernel = src1->type == GGML_TYPE_F32 && !y_non_contig; + + vk_matmul_pipeline mmp = ggml_vk_get_mul_mat_mat_id_pipeline(ctx, src0->type, y_non_contig ? GGML_TYPE_F16 : src1->type, (ggml_prec)dst->op_params[0]); + + const bool qx_needs_dequant = mmp == nullptr || x_non_contig; + const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !y_f32_kernel) || y_non_contig; + + if (qx_needs_dequant) { + GGML_ABORT("fatal error"); + } + + // Not implemented + GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT + + const uint64_t x_ne = ne01 * ne00; + const uint64_t y_ne = ne11 * ne10; + const uint64_t d_ne = ne21 * ne20; + + const uint32_t kpad = ggml_vk_align_size(ne10, ggml_vk_guess_matmul_id_pipeline_align(ctx, mmp, ne01, nei1)); + const bool aligned = ne10 == kpad && ne01 > 8 && nei1 > 8; + + vk_pipeline pipeline = ggml_vk_guess_matmul_id_pipeline(ctx, mmp, ne01, nei1, aligned); + + const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type); + const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); + const uint64_t x_sz = !qx_needs_dequant ? qx_sz : sizeof(ggml_fp16_t) * x_ne; + const uint64_t y_sz = y_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; + const uint64_t ids_sz = nbi2; + const uint64_t d_sz = sizeof(float) * d_ne; + + vk_pipeline to_fp16_vk_0 = nullptr; + vk_pipeline to_fp16_vk_1 = nullptr; + + if (x_non_contig) { + to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, GGML_TYPE_F16); + } else { + to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, src0->type); + } + if (y_non_contig) { + to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, GGML_TYPE_F16); + } else { + to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); + } + GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT + GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT + + if (dryrun) { + const uint64_t x_sz_upd = x_sz * ne02 * ne03; + const uint64_t y_sz_upd = y_sz * ne12 * ne13; + if ( + (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || + (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { + GGML_ABORT("Requested preallocation size is too large"); + } + if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { + ctx->prealloc_size_x = x_sz_upd; + } + if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { + ctx->prealloc_size_y = y_sz_upd; + } + + // Request descriptor sets + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + if (qx_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); + } + if (qy_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); + } + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + vk_buffer d_X; + uint64_t x_buf_offset = 0; + vk_buffer d_Y; + uint64_t y_buf_offset = 0; + if (!src0_uma) { + d_Qx = src0_buf_ctx->dev_buffer; + qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + if (!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qy != nullptr); + } + if (!ids_uma) { + d_ids = ids_buf_ctx->dev_buffer; + ids_buf_offset = vk_tensor_offset(ids) + ids->view_offs; + GGML_ASSERT(d_ids != nullptr); + } + if (qx_needs_dequant) { + d_X = ctx->prealloc_x; + GGML_ASSERT(d_X->size >= x_sz * ne02 * ne03); + } else { + d_X = d_Qx; + x_buf_offset = qx_buf_offset; + GGML_ASSERT(qx_sz == x_sz); + } + if (qy_needs_dequant) { + d_Y = ctx->prealloc_y; + GGML_ASSERT(d_Y->size >= y_sz * ne02 * ne03); + } else { + d_Y = d_Qy; + y_buf_offset = qy_buf_offset; + GGML_ASSERT(qy_sz == y_sz); + } + + if (x_non_contig) { + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); + } else if (qx_needs_dequant) { + const std::vector pc = { (uint32_t)ne01, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)(ggml_nelements(src0)) }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, + { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz * ne02 * ne03 }, vk_subbuffer{ d_X, 0, x_sz * ne02 * ne03 } }, pc.size() * sizeof(uint32_t), pc.data(), { (uint32_t)(x_ne * ne02 * ne03), 1, 1}); + } + if (y_non_contig) { + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); + } + + uint32_t stride_batch_x = ne00*ne01; + uint32_t stride_batch_y = ne10*ne11; + + if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { + stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); + } + + if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { + stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); + } + + // compute + ggml_vk_matmul_id( + ctx, subctx, pipeline, + { d_X, x_buf_offset, x_sz * ne02 * ne03 }, { d_Y, y_buf_offset, y_sz * ne12 * ne13 }, + { d_D, d_buf_offset, d_sz * ne22 * ne23 }, { d_ids, ids_buf_offset, ids_sz }, + ne01, ne21, ne10, ne10, ne10, ne01, + stride_batch_x, stride_batch_y, ne20*ne21, + n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11 + ); // NOLINT +} + +static void ggml_vk_mul_mat_vec_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_vec_id_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << ids << ", name=" << ids->name << ", type=" << ids->type << ", ne0=" << ids->ne[0] << ", ne1=" << ids->ne[1] << ", ne2=" << ids->ne[2] << ", ne3=" << ids->ne[3] << ", nb0=" << ids->nb[0] << ", nb1=" << ids->nb[1] << ", nb2=" << ids->nb[2] << ", nb3=" << ids->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); + GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT + GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT + GGML_ASSERT(ids->type == GGML_TYPE_I32); + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + const uint64_t ne03 = src0->ne[3]; + + const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + const uint64_t ne13 = src1->ne[3]; + + const uint64_t nei0 = ids->ne[0]; + const uint64_t nei1 = ids->ne[1]; + + const uint64_t nbi2 = ids->nb[2]; + + GGML_ASSERT(nei1 == 1); + + const uint64_t ne20 = dst->ne[0]; + const uint64_t ne21 = dst->ne[1]; + const uint64_t ne22 = dst->ne[2]; + const uint64_t ne23 = dst->ne[3]; + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + ggml_backend_vk_buffer_context * ids_buf_ctx = (ggml_backend_vk_buffer_context *)ids->buffer->context; + + vk_buffer d_Qx = nullptr; + size_t qx_buf_offset = 0; + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + vk_buffer d_ids = nullptr; + size_t ids_buf_offset = 0; + + bool src0_uma = false; + bool src1_uma = false; + bool ids_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + ggml_vk_host_get(ctx->device, ids->data, d_ids, ids_buf_offset); + src0_uma = d_Qx != nullptr; + src1_uma = d_Qy != nullptr; + ids_uma = d_ids != nullptr; + } + + const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); + const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); + + const bool f16_f32_kernel = src1->type == GGML_TYPE_F32; + + const bool qx_needs_dequant = x_non_contig; + const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !f16_f32_kernel) || y_non_contig; + + // Not implemented + GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT + + const uint64_t x_ne = ne01 * ne00; + const uint64_t y_ne = ne11 * ne10; + const uint64_t d_ne = ne21 * ne20; + + const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); + const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); + const uint64_t x_sz = x_non_contig ? ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment) : qx_sz; + const uint64_t y_sz = f16_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; + const uint64_t ids_sz = nbi2; + const uint64_t d_sz = sizeof(float) * d_ne; + + vk_pipeline to_fp16_vk_0 = nullptr; + vk_pipeline to_fp16_vk_1 = nullptr; + if (x_non_contig) { + to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, src0->type); + } + if (y_non_contig) { + to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, src1->type); + } else { + to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); + } + vk_pipeline dmmv = ggml_vk_get_dequantize_mul_mat_vec_id(ctx, src0->type, src1->type); + GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT + GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT + GGML_ASSERT(dmmv != nullptr); + + if (dryrun) { + const uint64_t x_sz_upd = x_sz * ne02 * ne03; + const uint64_t y_sz_upd = y_sz * ne12 * ne13; + if ( + (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || + (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { + GGML_ABORT("Requested preallocation size is too large"); + } + if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { + ctx->prealloc_size_x = x_sz_upd; + } + if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { + ctx->prealloc_size_y = y_sz_upd; + } + + // Request descriptor sets + if (qx_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); + } + if (qy_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); + } + ggml_pipeline_request_descriptor_sets(ctx->device, dmmv, 1); + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + vk_buffer d_X; + uint64_t x_buf_offset = 0; + vk_buffer d_Y; + uint64_t y_buf_offset = 0; + if(!src0_uma) { + d_Qx = src0_buf_ctx->dev_buffer; + qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + if(!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qy != nullptr); + } + if(!ids_uma) { + d_ids = ids_buf_ctx->dev_buffer; + ids_buf_offset = vk_tensor_offset(ids) + ids->view_offs; + GGML_ASSERT(d_ids != nullptr); + } + if (qx_needs_dequant) { + d_X = ctx->prealloc_x; + } else { + d_X = d_Qx; + x_buf_offset = qx_buf_offset; + GGML_ASSERT(qx_sz == x_sz); + } + if (qy_needs_dequant) { + d_Y = ctx->prealloc_y; + } else { + d_Y = d_Qy; + y_buf_offset = qy_buf_offset; + GGML_ASSERT(qy_sz == y_sz); + } + + if (x_non_contig) { + GGML_ASSERT(x_sz == ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment)); + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); + } + if (y_non_contig) { + GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); + } + + uint32_t stride_batch_y = ne10*ne11; + + if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { + stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); + } + + const uint32_t max_groups_x = ctx->device->properties.limits.maxComputeWorkGroupCount[0]; + + uint32_t groups_x = ne01; + uint32_t groups_z = 1; + + if (ne01 > max_groups_x) { + groups_z = 64; + groups_x = CEIL_DIV(groups_x, groups_z); + } + + // compute + const vk_mat_vec_id_push_constants pc = { + (uint32_t)ne00, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne01, + (uint32_t)x_ne, stride_batch_y, (uint32_t)(ne20*ne21), + (uint32_t)nei0, (uint32_t)ne11, + }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, dmmv, + { vk_subbuffer{ d_X, x_buf_offset, x_sz * ne02 * ne03 }, + vk_subbuffer{ d_Y, y_buf_offset, y_sz * ne12 * ne13 }, vk_subbuffer{ d_D, d_buf_offset, d_sz * ne22 * ne23}, vk_subbuffer{ d_ids, ids_buf_offset, ids_sz } }, + sizeof(vk_mat_vec_id_push_constants), &pc, { groups_x, (uint32_t)nei0, groups_z }); +} + +static void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_id(" << src0 << ", " << src1 << ", " << src2 << ", " << dst << ")"); + if (src2->ne[1] == 1 && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type))) { + ggml_vk_mul_mat_vec_id_q_f16(ctx, subctx, src0, src1, src2, dst, dryrun); + } else { + ggml_vk_mul_mat_id_q_f16(ctx, subctx, src0, src1, src2, dst, dryrun); + } +} + +static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_flash_attn((" << q << ", name=" << q->name << ", type=" << q->type << ", ne0=" << q->ne[0] << ", ne1=" << q->ne[1] << ", ne2=" << q->ne[2] << ", ne3=" << q->ne[3] << ", nb0=" << q->nb[0] << ", nb1=" << q->nb[1] << ", nb2=" << q->nb[2] << ", nb3=" << q->nb[3]; + std::cerr << "), (" << k << ", name=" << k->name << ", type=" << k->type << ", ne0=" << k->ne[0] << ", ne1=" << k->ne[1] << ", ne2=" << k->ne[2] << ", ne3=" << k->ne[3] << ", nb0=" << k->nb[0] << ", nb1=" << k->nb[1] << ", nb2=" << k->nb[2] << ", nb3=" << k->nb[3]; + std::cerr << "), (" << v << ", name=" << v->name << ", type=" << v->type << ", ne0=" << v->ne[0] << ", ne1=" << v->ne[1] << ", ne2=" << v->ne[2] << ", ne3=" << v->ne[3] << ", nb0=" << v->nb[0] << ", nb1=" << v->nb[1] << ", nb2=" << v->nb[2] << ", nb3=" << v->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne) + GGML_TENSOR_LOCALS(size_t, nbq, q, nb) + GGML_TENSOR_LOCALS(int64_t, nek, k, ne) + GGML_TENSOR_LOCALS(size_t, nbk, k, nb) + GGML_TENSOR_LOCALS(int64_t, nev, v, ne) + GGML_TENSOR_LOCALS(size_t, nbv, v, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + + const uint32_t nem1 = mask ? mask->ne[1] : 0; + const uint32_t nbm1 = mask ? mask->nb[1] : 0; + + const uint32_t D = neq0; + const uint32_t N = neq1; + const uint32_t KV = nek1; + + GGML_ASSERT(ne0 == D); + GGML_ASSERT(ne2 == N); + + // input tensor rows must be contiguous + GGML_ASSERT(nbq0 == ggml_type_size(q->type)); + GGML_ASSERT(nbk0 == ggml_type_size(k->type)); + GGML_ASSERT(nbv0 == ggml_type_size(v->type)); + + GGML_ASSERT(neq0 == D); + GGML_ASSERT(nek0 == D); + GGML_ASSERT(nev0 == D); + + GGML_ASSERT(neq1 == N); + GGML_ASSERT(nev0 == D); + + GGML_ASSERT(nev1 == nek1); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + assert(dst->type == GGML_TYPE_F32); + assert(q->type == GGML_TYPE_F32); + assert(k->type == v->type); + + vk_pipeline *pipelines; + // XXX TODO other backends may be changing accumulator precision to default to f32 soon + bool f32acc = dst->op_params[3] == GGML_PREC_F32; + bool small_rows = N <= flash_attention_num_small_rows; + switch (D) { + case 64: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D64[k->type][f32acc][small_rows][0]; break; + case 80: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D80[k->type][f32acc][small_rows][0]; break; + case 96: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D96[k->type][f32acc][small_rows][0]; break; + case 112: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D112[k->type][f32acc][small_rows][0]; break; + case 128: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D128[k->type][f32acc][small_rows][0]; break; + case 256: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D256[k->type][f32acc][small_rows][0]; break; + default: + assert(!"unsupported D value"); + return; + } + assert(pipelines); + + bool aligned = (KV % pipelines[1]->align) == 0; + vk_pipeline pipeline = pipelines[aligned]; + assert(pipeline); + + if (dryrun) { + // Request descriptor sets + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + return; + } + + float scale = 1.0f; + float max_bias = 0.0f; + float logit_softcap = 0.0f; + + memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + + if (logit_softcap != 0) { + scale /= logit_softcap; + } + + const uint32_t n_head_kv = neq2; + const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head_kv)); + const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); + const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + + ggml_vk_sync_buffers(subctx); + + vk_buffer d_Q = nullptr, d_K = nullptr, d_V = nullptr, d_D = nullptr, d_M = nullptr; + size_t q_buf_offset = 0, k_buf_offset = 0, v_buf_offset = 0, d_buf_offset = 0, m_buf_offset = 0; + + bool Q_uma = false, K_uma = false, V_uma = false, D_uma = false, M_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, q->data, d_Q, q_buf_offset); + ggml_vk_host_get(ctx->device, k->data, d_K, q_buf_offset); + ggml_vk_host_get(ctx->device, v->data, d_V, q_buf_offset); + ggml_vk_host_get(ctx->device, dst->data, d_D, q_buf_offset); + Q_uma = d_Q != nullptr; + K_uma = d_K != nullptr; + V_uma = d_V != nullptr; + D_uma = d_D != nullptr; + if (mask) { + ggml_vk_host_get(ctx->device, mask->data, d_M, q_buf_offset); + M_uma = d_M != nullptr; + } + } + + + ggml_backend_vk_buffer_context * d_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * q_buf_ctx = (ggml_backend_vk_buffer_context *)q->buffer->context; + ggml_backend_vk_buffer_context * k_buf_ctx = (ggml_backend_vk_buffer_context *)k->buffer->context; + ggml_backend_vk_buffer_context * v_buf_ctx = (ggml_backend_vk_buffer_context *)v->buffer->context; + + if (!Q_uma) { + d_Q = q_buf_ctx->dev_buffer; + q_buf_offset = vk_tensor_offset(q) + q->view_offs; + } + if (!K_uma) { + d_K = k_buf_ctx->dev_buffer; + k_buf_offset = vk_tensor_offset(k) + k->view_offs; + } + if (!V_uma) { + d_V = v_buf_ctx->dev_buffer; + v_buf_offset = vk_tensor_offset(v) + v->view_offs; + } + if (!D_uma) { + d_D = d_buf_ctx->dev_buffer; + d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + } + + if (!M_uma) { + d_M = d_Q; + m_buf_offset = q_buf_offset; + if (mask) { + ggml_backend_vk_buffer_context * m_buf_ctx = (ggml_backend_vk_buffer_context*)mask->buffer->context; + d_M = m_buf_ctx->dev_buffer; + m_buf_offset = vk_tensor_offset(mask) + mask->view_offs; + } + } + + const vk_flash_attn_push_constants pc = { N, KV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3, (uint32_t)neq2, (uint32_t)neq3, (uint32_t)nek2, (uint32_t)nek3, (uint32_t)nev2, (uint32_t)nev3, nem1, (uint32_t)nbq2, (uint32_t)nbq3, (uint32_t)nbk2, (uint32_t)nbk3, (uint32_t)nbv2, (uint32_t)nbv3, nbm1, scale, max_bias, logit_softcap, mask != nullptr, n_head_log2, m0, m1 }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + vk_subbuffer{d_Q, q_buf_offset, VK_WHOLE_SIZE}, + vk_subbuffer{d_K, k_buf_offset, VK_WHOLE_SIZE}, + vk_subbuffer{d_V, v_buf_offset, VK_WHOLE_SIZE}, + vk_subbuffer{d_M, m_buf_offset, VK_WHOLE_SIZE}, + vk_subbuffer{d_D, d_buf_offset, VK_WHOLE_SIZE}, + }, + sizeof(vk_flash_attn_push_constants), &pc, { (uint32_t)neq1, (uint32_t)neq2, (uint32_t)neq3 }); +} + +static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, ggml_op op) { + switch (op) { + case GGML_OP_GET_ROWS: + GGML_ASSERT(src1->type == GGML_TYPE_I32); + if (dst->type == GGML_TYPE_F16) { + return ctx->device->pipeline_get_rows[src0->type]; + } + if (dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_get_rows_f32[src0->type]; + } + return nullptr; + case GGML_OP_ACC: + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_acc_f32; + } + return nullptr; + case GGML_OP_ADD: + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_add_f32_norepeat : ctx->device->pipeline_add_f32; + } + if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) { + return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_add_f16_f32_f16_norepeat : ctx->device->pipeline_add_f16_f32_f16; + } + return nullptr; + case GGML_OP_MUL: + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_mul_f32_norepeat : ctx->device->pipeline_mul_f32; + } + return nullptr; + case GGML_OP_DIV: + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_div_f32_norepeat : ctx->device->pipeline_div_f32; + } + return nullptr; + case GGML_OP_CONCAT: + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_concat_f32; + } + if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { + return ctx->device->pipeline_concat_f16; + } + if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_I32 && dst->type == GGML_TYPE_I32) { + return ctx->device->pipeline_concat_i32; + } + return nullptr; + case GGML_OP_UPSCALE: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_upscale_f32; + } + return nullptr; + case GGML_OP_SCALE: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_scale_f32; + } + return nullptr; + case GGML_OP_SQR: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_sqr_f32; + } + return nullptr; + case GGML_OP_SIN: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_sin_f32; + } + return nullptr; + case GGML_OP_COS: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_cos_f32; + } + return nullptr; + case GGML_OP_CLAMP: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_clamp_f32; + } + return nullptr; + case GGML_OP_PAD: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_pad_f32; + } + return nullptr; + case GGML_OP_REPEAT: + if (ggml_type_size(src0->type) == sizeof(float) && ggml_type_size(dst->type) == sizeof(float)) { + return ctx->device->pipeline_repeat_f32; + } + return nullptr; + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_DUP: + return ggml_vk_get_cpy_pipeline(ctx, src0, dst, dst->type); + case GGML_OP_NORM: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_norm_f32; + } + return nullptr; + case GGML_OP_GROUP_NORM: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_group_norm_f32; + } + return nullptr; + case GGML_OP_RMS_NORM: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_rms_norm_f32; + } + return nullptr; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(dst)) { + case GGML_UNARY_OP_SILU: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_silu_f32; + } + break; + case GGML_UNARY_OP_GELU: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_gelu_f32; + } + break; + case GGML_UNARY_OP_GELU_QUICK: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_gelu_quick_f32; + } + break; + case GGML_UNARY_OP_RELU: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_relu_f32; + } + break; + case GGML_UNARY_OP_TANH: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_tanh_f32; + } + break; + default: + break; + } + return nullptr; + case GGML_OP_DIAG_MASK_INF: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_diag_mask_inf_f32; + } + return nullptr; + case GGML_OP_SOFT_MAX: + GGML_ASSERT(!src1 || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); + + if (src0->type == GGML_TYPE_F32 && (src1 == nullptr || src1->type == GGML_TYPE_F32) && dst->type == GGML_TYPE_F32) { + return src0->ne[0] > 1024 ? ctx->device->pipeline_soft_max_f32_wg512 : ctx->device->pipeline_soft_max_f32; + } + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) { + return src0->ne[0] > 1024 ? ctx->device->pipeline_soft_max_f32_f16_wg512 : ctx->device->pipeline_soft_max_f32_f16; + } + return nullptr; + case GGML_OP_ROPE: + { + const int mode = ((const int32_t *) dst->op_params)[2]; + const bool is_neox = mode & GGML_ROPE_TYPE_NEOX; + + if (is_neox) { + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_rope_neox_f32; + } + if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { + return ctx->device->pipeline_rope_neox_f16; + } + } else { + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_rope_norm_f32; + } + if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { + return ctx->device->pipeline_rope_norm_f16; + } + } + return nullptr; + } + case GGML_OP_ARGSORT: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_I32) { + return ctx->device->pipeline_argsort_f32; + } + return nullptr; + case GGML_OP_SUM_ROWS: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_sum_rows_f32; + } + return nullptr; + case GGML_OP_IM2COL: + if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_im2col_f32; + } + if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) { + return ctx->device->pipeline_im2col_f32_f16; + } + return nullptr; + case GGML_OP_TIMESTEP_EMBEDDING: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_timestep_embedding_f32; + } + return nullptr; + case GGML_OP_POOL_2D: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_pool2d_f32; + } + return nullptr; + case GGML_OP_RWKV_WKV6: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_rwkv_wkv6_f32; + } + return nullptr; + case GGML_OP_LEAKY_RELU: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_leaky_relu_f32; + } + return nullptr; + default: + return nullptr; + } + + GGML_UNUSED(src2); +} + +static bool ggml_vk_op_supports_incontiguous(ggml_op op) { + switch (op) { + case GGML_OP_CPY: + case GGML_OP_GET_ROWS: + case GGML_OP_ADD: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_REPEAT: + return true; + default: + return false; + } +} + +static uint32_t get_misalign_bytes(ggml_backend_vk_context * ctx, const ggml_tensor * t) +{ + return ((vk_tensor_offset(t) + t->view_offs) & (ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1));; +} + +template void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, T &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { + GGML_UNUSED(p); + GGML_UNUSED(src0); + GGML_UNUSED(src1); + GGML_UNUSED(src2); + GGML_UNUSED(dst); + static_assert(!std::is_const::value, "unexpected type"); + GGML_ASSERT(!src0 || get_misalign_bytes(ctx, src0) == 0); + GGML_ASSERT(!src1 || get_misalign_bytes(ctx, src1) == 0); + GGML_ASSERT(!src2 || get_misalign_bytes(ctx, src2) == 0); + GGML_ASSERT(!dst || get_misalign_bytes(ctx, dst) == 0); +} + +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_unary_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { + const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + p.misalign_offsets = (a_offset << 16) | d_offset; + + GGML_UNUSED(src1); + GGML_UNUSED(src2); +} + +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_binary_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { + const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + const uint32_t b_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); + const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + GGML_ASSERT(dst->op != GGML_OP_GET_ROWS || (a_offset == 0 && b_offset == 0 && d_offset == 0)); + + p.misalign_offsets = (a_offset << 16) | (b_offset << 8) | d_offset; + + GGML_UNUSED(src2); +} + +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_upscale_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { + const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + p.a_offset = a_offset; + p.d_offset = d_offset; + + GGML_UNUSED(src1); + GGML_UNUSED(src2); +} + +template +static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, ggml_op op, PC&& pc, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_op_f32((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + if (src1 != nullptr) { + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + } + if (src2 != nullptr) { + std::cerr << "), (" << src2 << ", name=" << src2->name << ", type=" << src2->type << ", ne0=" << src2->ne[0] << ", ne1=" << src2->ne[1] << ", ne2=" << src2->ne[2] << ", ne3=" << src2->ne[3] << ", nb0=" << src2->nb[0] << ", nb1=" << src2->nb[1] << ", nb2=" << src2->nb[2] << ", nb3=" << src2->nb[3]; + } + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << ggml_op_name(op) << ", " << (dryrun ? "dryrun" : "") << ")"); + GGML_ASSERT(op == GGML_OP_GET_ROWS || (!ggml_is_quantized(src0->type) && (src1 == nullptr || !ggml_is_quantized(src1->type)))); // NOLINT + GGML_ASSERT(ggml_vk_op_supports_incontiguous(op) || ggml_vk_dim01_contiguous(src0)); // NOLINT + GGML_ASSERT(dst->buffer != nullptr); + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + const uint64_t ne03 = src0->ne[3]; + const uint64_t ne0 = ne00 * ne01; + + const bool use_src1 = src1 != nullptr; + const uint64_t ne10 = use_src1 ? src1->ne[0] : 0; + const uint64_t ne11 = use_src1 ? src1->ne[1] : 0; + const uint64_t ne12 = use_src1 ? src1->ne[2] : 0; + const uint64_t ne13 = use_src1 ? src1->ne[3] : 0; + const uint64_t ne1 = ne10 * ne11; + // const uint64_t nb10 = use_src1 ? src1->nb[0] : 0; + + const bool use_src2 = src2 != nullptr; + const uint64_t ne20 = use_src2 ? src2->ne[0] : 0; + const uint64_t ne21 = use_src2 ? src2->ne[1] : 0; + const uint64_t ne22 = use_src2 ? src2->ne[2] : 0; + const uint64_t ne23 = use_src2 ? src2->ne[3] : 0; + const uint64_t ne2 = ne20 * ne21; + + const uint64_t ned0 = dst->ne[0]; + const uint64_t ned1 = dst->ne[1]; + const uint64_t ned2 = dst->ne[2]; + const uint64_t ned3 = dst->ne[3]; + const uint64_t ned = ned0 * ned1; + + init_pushconst_fastdiv(pc); + + vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, src2, dst, op); + + if (pipeline == nullptr) { + std::cerr << "ggml_vulkan: Error: Missing op: " << ggml_op_name(op) << " for " << ggml_type_name(src0->type); + if (src1 != nullptr) { + std::cerr << " and " << ggml_type_name(src1->type); + } + std::cerr << " to " << ggml_type_name(dst->type) << std::endl; + GGML_ABORT("fatal error"); + } + + if (dryrun) { + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + return; + } + + const bool op_supports_incontiguous = ggml_vk_op_supports_incontiguous(op); + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = use_src1 ? (ggml_backend_vk_buffer_context *)src1->buffer->context : nullptr; + ggml_backend_vk_buffer_context * src2_buf_ctx = use_src2 ? (ggml_backend_vk_buffer_context *)src2->buffer->context : nullptr; + + vk_buffer d_X = nullptr; + size_t x_buf_offset = 0; + vk_buffer d_Y = nullptr; + size_t y_buf_offset = 0; + vk_buffer d_Z = nullptr; + size_t z_buf_offset = 0; + + bool src0_uma = false; + bool src1_uma = false; + bool src2_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src0->data, d_X, x_buf_offset); + src0_uma = d_X != nullptr; + if (use_src1) { + ggml_vk_host_get(ctx->device, src1->data, d_Y, y_buf_offset); + src1_uma = d_Y != nullptr; + } + if (use_src2) { + ggml_vk_host_get(ctx->device, src2->data, d_Z, z_buf_offset); + src2_uma = d_Z != nullptr; + } + } + + uint64_t x_sz = ggml_type_size(src0->type)/ggml_blck_size(src0->type) * ne0; + uint64_t y_sz = use_src1 ? ggml_type_size(src1->type) * ne1 : 0; + uint64_t z_sz = use_src2 ? ggml_type_size(src2->type) * ne2 : 0; + uint64_t d_sz = ggml_type_size(dst->type) * ned; + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + + // Workaround for tiny tensor inputs on ROPE + if (op == GGML_OP_ROPE && use_src1 && y_sz > d_D->size) { + y_sz = VK_WHOLE_SIZE; + } + + GGML_ASSERT(d_D != nullptr); + uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + if(!src0_uma) { + d_X = src0_buf_ctx->dev_buffer; + x_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_X != nullptr); + } + if (use_src1 && !src1_uma) { + d_Y = src1_buf_ctx->dev_buffer; + y_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Y != nullptr); + } + if (use_src2 && !src2_uma) { + d_Z = src2_buf_ctx->dev_buffer; + z_buf_offset = vk_tensor_offset(src2) + src2->view_offs; + GGML_ASSERT(d_Z != nullptr); + } + // Compute misalignment offset for descriptors and store it in in push constants, then align the descriptor offsets. + init_pushconst_tensor_offsets(ctx, pc, src0, src1, src2, dst); + x_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); + y_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); + z_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); + d_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); + + if (op_supports_incontiguous) { + x_sz = ggml_nbytes(src0); + y_sz = use_src1 ? ggml_nbytes(src1) : 0; + z_sz = use_src2 ? ggml_nbytes(src2) : 0; + d_sz = ggml_nbytes(dst); + + if (x_buf_offset + x_sz >= d_X->size) { + x_sz = VK_WHOLE_SIZE; + } + if (use_src1 && y_buf_offset + y_sz >= d_Y->size) { + y_sz = VK_WHOLE_SIZE; + } + if (use_src2 && z_buf_offset + z_sz >= d_Z->size) { + z_sz = VK_WHOLE_SIZE; + } + if (d_buf_offset + d_sz >= d_D->size) { + d_sz = VK_WHOLE_SIZE; + } + } + + std::array elements; + + // Single call if dimension 2 is contiguous + GGML_ASSERT(op_supports_incontiguous || (ggml_is_contiguous(src0) && (src1 == nullptr || ggml_is_contiguous(src1)))); + + switch (op) { + case GGML_OP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_SOFT_MAX: + case GGML_OP_SUM_ROWS: + { + const uint32_t nr = ggml_nrows(src0); + if (nr > 262144) { + elements = { 512, 512, CEIL_DIV(nr, 262144) }; + } else if (nr > 512) { + elements = { 512, CEIL_DIV(nr, 512), 1 }; + } else { + elements = { nr, 1, 1 }; + } + } break; + case GGML_OP_GROUP_NORM: + { + const uint32_t num_groups = dst->op_params[0]; + elements = { num_groups * (uint32_t)src0->ne[3], 1, 1 }; + } break; + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_ROPE: + elements = { (uint32_t)ggml_nrows(src0), (uint32_t)ne00, 1 }; + break; + case GGML_OP_GET_ROWS: + elements = { (uint32_t)ne00, (uint32_t)ne10, (uint32_t)(ne11 * ne12) }; + break; + case GGML_OP_ARGSORT: + elements = { (uint32_t)ne00, (uint32_t)ggml_nrows(src0), 1 }; + break; + case GGML_OP_IM2COL: + { + const bool is_2D = dst->op_params[6] == 1; + + const uint32_t IC = src1->ne[is_2D ? 2 : 1]; + + const uint32_t KH = is_2D ? src0->ne[1] : 1; + const uint32_t KW = src0->ne[0]; + + const uint32_t OH = is_2D ? dst->ne[2] : 1; + const uint32_t OW = dst->ne[1]; + + const uint32_t batch = src1->ne[is_2D ? 3 : 2]; + + elements = { OW * KW * KH, OH, batch * IC }; + } break; + case GGML_OP_TIMESTEP_EMBEDDING: + { + const uint32_t dim = dst->op_params[0]; + uint32_t half_ceil = (dim + 1) / 2; + elements = { half_ceil, (uint32_t)src0->ne[0], 1 }; + } break; + case GGML_OP_POOL_2D: + { + const uint32_t N = dst->ne[3]; + const uint32_t OC = dst->ne[2]; + const uint32_t OH = dst->ne[1]; + const uint32_t OW = dst->ne[0]; + elements = { N * OC * OH * OW, 1, 1}; + } break; + case GGML_OP_ADD: + case GGML_OP_DIV: + case GGML_OP_MUL: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_REPEAT: + case GGML_OP_CPY: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_UNARY: + { + const uint32_t ne = ggml_nelements(dst); + if (ne > 262144) { + elements = { 512, 512, CEIL_DIV(ne, 262144) }; + } else if (ne > 512) { + elements = { 512, CEIL_DIV(ne, 512), 1 }; + } else { + elements = { ne, 1, 1 }; + } + } break; + default: + elements = { (uint32_t)ggml_nelements(src0), 1, 1 }; + break; + } + + if (!op_supports_incontiguous) { + if (x_sz != VK_WHOLE_SIZE) { + x_sz *= ne02 * ne03; + } + if (use_src1 && y_sz != VK_WHOLE_SIZE) { + y_sz *= ne12 * ne13; + } + if (use_src2 && z_sz != VK_WHOLE_SIZE) { + z_sz *= ne22 * ne23; + } + if (d_sz != VK_WHOLE_SIZE) { + d_sz *= ned2 * ned3; + } + } + + if (op == GGML_OP_SOFT_MAX) { + // Empty src1 is possible in soft_max, but the shader needs a buffer + vk_subbuffer subbuf_y; + if (use_src1) { + subbuf_y = { d_Y, y_buf_offset, y_sz }; + } else { + subbuf_y = { d_X, 0, x_sz }; + } + + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, subbuf_y, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } else if (op == GGML_OP_ROPE) { + // Empty src2 is possible in rope, but the shader needs a buffer + vk_subbuffer subbuf_z; + if (use_src2) { + subbuf_z = { d_Z, z_buf_offset, z_sz }; + } else { + subbuf_z = { d_X, 0, x_sz }; + } + + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, subbuf_z, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } else if (op == GGML_OP_IM2COL) { + // im2col uses only src1 and dst buffers + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } else if (use_src2) { + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_Z, z_buf_offset, z_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } else if (use_src1) { + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } else { + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } +} + +static void ggml_vk_get_rows(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_GET_ROWS, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, 0, + }, dryrun); +} + +static void ggml_vk_acc(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + int nb1 = dst->op_params[0] / 4; // 4 bytes of float32 + int nb2 = dst->op_params[1] / 4; // 4 bytes of float32 + // int nb3 = dst->op_params[2] / 4; // 4 bytes of float32 - unused + int offset = dst->op_params[3] / 4; // offset in bytes + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_ACC, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)nb1, (uint32_t)nb2, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t)nb1, (uint32_t)nb2, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, offset, + }, dryrun); +} + +static void ggml_vk_add(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_ADD, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, 0, + }, dryrun); +} + +static void ggml_vk_mul(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_MUL, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, 0, + }, dryrun); +} + +static void ggml_vk_div(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_DIV, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, 0, + }, dryrun); +} + +static void ggml_vk_op_f32_rwkv6(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, const vk_op_rwkv_wkv6_push_constants&& pc, bool dryrun = false) { + const ggml_tensor * k = dst->src[0]; + const ggml_tensor * v = dst->src[1]; + const ggml_tensor * r = dst->src[2]; + const ggml_tensor * tf = dst->src[3]; + const ggml_tensor * td = dst->src[4]; + const ggml_tensor * state = dst->src[5]; + + GGML_ASSERT(!ggml_is_quantized(k->type)); + GGML_ASSERT(!ggml_is_quantized(v->type)); + GGML_ASSERT(!ggml_is_quantized(r->type)); + GGML_ASSERT(!ggml_is_quantized(tf->type)); + GGML_ASSERT(!ggml_is_quantized(td->type)); + GGML_ASSERT(!ggml_is_quantized(state->type)); + GGML_ASSERT(dst->buffer != nullptr); + + vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, k, v, r, dst, GGML_OP_RWKV_WKV6); + GGML_ASSERT(pipeline != nullptr); + + if (dryrun) { + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + return; + } + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * k_buf_ctx = (ggml_backend_vk_buffer_context *)k->buffer->context; + ggml_backend_vk_buffer_context * v_buf_ctx = (ggml_backend_vk_buffer_context *)v->buffer->context; + ggml_backend_vk_buffer_context * r_buf_ctx = (ggml_backend_vk_buffer_context *)r->buffer->context; + ggml_backend_vk_buffer_context * tf_buf_ctx = (ggml_backend_vk_buffer_context *)tf->buffer->context; + ggml_backend_vk_buffer_context * td_buf_ctx = (ggml_backend_vk_buffer_context *)td->buffer->context; + ggml_backend_vk_buffer_context * state_buf_ctx = (ggml_backend_vk_buffer_context *)state->buffer->context; + + ggml_vk_sync_buffers(subctx); + + vk_buffer d_D = nullptr, d_K = nullptr, d_V = nullptr, d_R = nullptr, d_TF = nullptr, d_TD = nullptr, d_State = nullptr; + size_t k_offset = 0, v_offset = 0, r_offset = 0, tf_offset = 0, td_offset = 0, state_offset = 0, dst_offset = 0; + bool K_uma = false, V_uma = false, R_uma = false, TF_uma = false, TD_uma = false, STATE_uma = false, DST_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, k->data, d_K, k_offset); + ggml_vk_host_get(ctx->device, v->data, d_V, v_offset); + ggml_vk_host_get(ctx->device, r->data, d_R, r_offset); + ggml_vk_host_get(ctx->device, tf->data, d_TF, tf_offset); + ggml_vk_host_get(ctx->device, td->data, d_TD, td_offset); + ggml_vk_host_get(ctx->device, state->data, d_State, state_offset); + ggml_vk_host_get(ctx->device, dst->data, d_D, dst_offset); + + K_uma = d_K != nullptr; + V_uma = d_V != nullptr; + R_uma = d_R != nullptr; + TF_uma = d_TF != nullptr; + TD_uma = d_TD != nullptr; + STATE_uma = d_State != nullptr; + DST_uma = d_D != nullptr; + } + + if (!K_uma) { + d_K = k_buf_ctx->dev_buffer; + k_offset = vk_tensor_offset(k) + k->view_offs; + } + if (!V_uma) { + d_V = v_buf_ctx->dev_buffer; + v_offset = vk_tensor_offset(v) + v->view_offs; + } + if (!R_uma) { + d_R = r_buf_ctx->dev_buffer; + r_offset = vk_tensor_offset(r) + r->view_offs; + } + if (!TF_uma) { + d_TF = tf_buf_ctx->dev_buffer; + tf_offset = vk_tensor_offset(tf) + tf->view_offs; + } + if (!TD_uma) { + d_TD = td_buf_ctx->dev_buffer; + td_offset = vk_tensor_offset(td) + td->view_offs; + } + if (!STATE_uma) { + d_State = state_buf_ctx->dev_buffer; + state_offset = vk_tensor_offset(state) + state->view_offs; + } + if (!DST_uma) { + d_D = dst_buf_ctx->dev_buffer; + dst_offset = vk_tensor_offset(dst) + dst->view_offs; + } + + const uint64_t k_size = ggml_nbytes(k); + const uint64_t v_size = ggml_nbytes(v); + const uint64_t r_size = ggml_nbytes(r); + const uint64_t tf_size = ggml_nbytes(tf); + const uint64_t td_size = ggml_nbytes(td); + const uint64_t state_size = ggml_nbytes(state); + const uint64_t dst_size = ggml_nbytes(dst); + + std::array elements = { + (uint32_t)(pc.B * pc.H), + 1, + 1 + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { + vk_subbuffer{ d_K, k_offset, k_size }, + vk_subbuffer{ d_V, v_offset, v_size }, + vk_subbuffer{ d_R, r_offset, r_size }, + vk_subbuffer{ d_TF, tf_offset, tf_size }, + vk_subbuffer{ d_TD, td_offset, td_size }, + vk_subbuffer{ d_State, state_offset, state_size }, + vk_subbuffer{ d_D, dst_offset, dst_size } + }, sizeof(vk_op_rwkv_wkv6_push_constants), &pc, elements); +} + +static void ggml_vk_rwkv_wkv6(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, bool dryrun = false) { + const size_t seq_length = dst->src[0]->ne[3]; + const size_t n_embed = dst->ne[0]; + const size_t n_heads = dst->src[0]->ne[2]; + const size_t n_seqs = dst->src[5]->ne[1]; + + ggml_vk_op_f32_rwkv6( + ctx, subctx, dst, + { + (uint32_t)n_seqs, + (uint32_t)seq_length, + (uint32_t)n_embed, + (uint32_t)n_heads, + }, + dryrun + ); +} + +static void ggml_vk_concat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + int * op_params = (int *)dst->op_params; + + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_CONCAT, { + (uint32_t)ggml_nelements(dst), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, op_params[0], + }, dryrun); +} + +static void ggml_vk_upscale(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + + const float sf0 = (float)dst->ne[0] / src0->ne[0]; + const float sf1 = (float)dst->ne[1] / src0->ne[1]; + const float sf2 = (float)dst->ne[2] / src0->ne[2]; + const float sf3 = (float)dst->ne[3] / src0->ne[3]; + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_UPSCALE, { + (uint32_t)ggml_nelements(dst), 0, 0, + (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)dst->ne[0], (uint32_t)dst->ne[1], (uint32_t)dst->ne[2],(uint32_t)dst->ne[3], + sf0, sf1, sf2, sf3, + }, dryrun); +} + +static void ggml_vk_scale(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + float * op_params = (float *)dst->op_params; + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SCALE, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + op_params[0], 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_sqr(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SQR, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_sin(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SIN, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_cos(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_COS, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_clamp(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + float * op_params = (float *)dst->op_params; + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_CLAMP, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + op_params[0], op_params[1], + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_pad(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_PAD, { + (uint32_t)ggml_nelements(dst), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_repeat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_REPEAT, { + (uint32_t)ggml_nelements(dst), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_cpy(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_CPY, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + float * op_params = (float *)dst->op_params; + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_NORM, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0], 0.0f }, dryrun); +} + +static void ggml_vk_group_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const int * int_op_params = (const int *)dst->op_params; + const float * float_op_params = (const float *)dst->op_params; + + const uint32_t num_groups = int_op_params[0]; + const float eps = float_op_params[1]; + const uint32_t group_size = src0->ne[0] * src0->ne[1] * ((src0->ne[2] + num_groups - 1) / num_groups); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_GROUP_NORM, { group_size, 0, eps, 0.0f }, dryrun); +} + +static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + float * op_params = (float *)dst->op_params; + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_RMS_NORM, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0], 0.0f }, dryrun); +} + +static void ggml_vk_unary(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_UNARY, { (uint32_t)ggml_nelements(src0), 0, 0.0f, 0.0f }, dryrun); +} + +static void ggml_vk_diag_mask_inf(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + int32_t * op_params = (int32_t *)dst->op_params; + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_DIAG_MASK_INF, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0] }, dryrun); +} + +static void ggml_vk_soft_max(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + float * op_params = (float *)dst->op_params; + + float scale = op_params[0]; + float max_bias = op_params[1]; + + const uint32_t ncols = (uint32_t)src0->ne[0]; + const uint32_t nrows_x = (uint32_t)ggml_nrows(src0); + const uint32_t nrows_y = (uint32_t)src0->ne[1]; + + const uint32_t n_head_kv = nrows_x/nrows_y; + const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head_kv)); + + const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); + const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_SOFT_MAX, { + ncols, + src1 != nullptr ? nrows_y : (uint32_t)0, + scale, max_bias, + m0, m1, + n_head_log2, + nrows_x, + }, dryrun); +} + +static void ggml_vk_rope(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, bool dryrun = false) { + const int n_dims = ((int32_t *) dst->op_params)[1]; + // const int mode = ((int32_t *) dst->op_params)[2]; + // const int n_ctx = ((int32_t *) dst->op_params)[3]; + const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; + const float freq_base = ((float *) dst->op_params)[5]; + const float freq_scale = ((float *) dst->op_params)[6]; + const float ext_factor = ((float *) dst->op_params)[7]; + const float attn_factor = ((float *) dst->op_params)[8]; + const float beta_fast = ((float *) dst->op_params)[9]; + const float beta_slow = ((float *) dst->op_params)[10]; + + float corr_dims[2]; + ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); + + const float theta_scale = powf(freq_base, -2.0f/n_dims); + + ggml_vk_op_f32(ctx, subctx, src0, src1, src2, dst, GGML_OP_ROPE, { + (uint32_t)src0->ne[0], (uint32_t)n_dims, freq_scale, (uint32_t)src0->ne[1], + freq_base, ext_factor, attn_factor, {corr_dims[0], corr_dims[1]}, theta_scale, + src2 != nullptr, + }, dryrun); +} + +static void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + int32_t * op_params = (int32_t *)dst->op_params; + + uint32_t ncols = src0->ne[0]; + + uint32_t ncols_pad = 1; + while (ncols_pad < ncols) { + ncols_pad *= 2; + } + + GGML_ASSERT(ncols_pad <= 1024); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_ARGSORT, { + ncols, + ncols_pad, + op_params[0], + }, dryrun); +} + +static void ggml_vk_sum_rows(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SUM_ROWS, { (uint32_t)src0->ne[0], 0, 0.0f, 0.0f }, dryrun); +} + +static void ggml_vk_im2col(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const int32_t s0 = dst->op_params[0]; + const int32_t s1 = dst->op_params[1]; + const int32_t p0 = dst->op_params[2]; + const int32_t p1 = dst->op_params[3]; + const int32_t d0 = dst->op_params[4]; + const int32_t d1 = dst->op_params[5]; + + const bool is_2D = dst->op_params[6] == 1; + + const uint32_t IC = src1->ne[is_2D ? 2 : 1]; + const uint32_t IH = is_2D ? src1->ne[1] : 1; + const uint32_t IW = src1->ne[0]; + + const uint32_t KH = is_2D ? src0->ne[1] : 1; + const uint32_t KW = src0->ne[0]; + + const uint32_t OH = is_2D ? dst->ne[2] : 1; + const uint32_t OW = dst->ne[1]; + + const uint32_t offset_delta = src1->nb[is_2D ? 2 : 1] / 4; // nb is byte offset, src is type float32 + const uint32_t batch_offset = src1->nb[is_2D ? 3 : 2] / 4; // nb is byte offset, src is type float32 + + const uint32_t pelements = OW * KW * KH; + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_IM2COL, { + batch_offset, offset_delta, + IC, IW, IH, OW, OH, KW, KH, + pelements, + IC * KH * KW, + s0, s1, p0, p1, d0, d1, + }, dryrun); +} + +static void ggml_vk_timestep_embedding(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t dim = dst->op_params[0]; + const uint32_t max_period = dst->op_params[1]; + const uint32_t nb1 = dst->nb[1] / ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_TIMESTEP_EMBEDDING, { + nb1, dim, max_period, + }, dryrun); +} + +static void ggml_vk_pool_2d(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + uint32_t op = static_cast(dst->op_params[0]); + const int32_t k1 = dst->op_params[1]; + const int32_t k0 = dst->op_params[2]; + const int32_t s1 = dst->op_params[3]; + const int32_t s0 = dst->op_params[4]; + const int32_t p1 = dst->op_params[5]; + const int32_t p0 = dst->op_params[6]; + + const uint32_t IH = src0->ne[1]; + const uint32_t IW = src0->ne[0]; + + const uint32_t N = dst->ne[3]; + + const uint32_t OC = dst->ne[2]; + const uint32_t OH = dst->ne[1]; + const uint32_t OW = dst->ne[0]; + + const uint32_t parallel_elements = N * OC * OH * OW; + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_POOL_2D, { + IW, IH, OW, OH, OC, + parallel_elements, + op, + k0, k1, s0, s1, p0, p1, + }, dryrun); +} + +static void ggml_vk_leaky_relu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const float * op_params = (const float *)dst->op_params; + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_LEAKY_RELU, { (uint32_t)ggml_nelements(src0), 0, op_params[0], 0.0f }, dryrun); +} + +#ifdef GGML_VULKAN_RUN_TESTS +static void ggml_vk_print_matrix_area(const void * data, ggml_type type, int ne0, int ne1, int i0, int i1, int i2) { + if (type != GGML_TYPE_F32 && type != GGML_TYPE_F16) { + return; + } + i0 = std::max(i0, 5); + i1 = std::max(i1, 5); + i2 = std::max(i2, 0); + fprintf(stderr, " "); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + fprintf(stderr, "%7d ", idx1); + } + fprintf(stderr, "\n"); + for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { + fprintf(stderr, "%7d: ", idx0); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + if (idx0 >= 0 && idx0 < ne0 && idx1 >= 0 && idx1 < ne1) { + float val; + if (type == GGML_TYPE_F32) { + val = *((const float *) data + i2*ne1*ne0 + idx1*ne0 + idx0); + } else if (type == GGML_TYPE_F16) { + val = ggml_fp16_to_fp32(*((const ggml_fp16_t *) data + i2*ne1*ne0 + idx1*ne0 + idx0)); + } else { + GGML_ABORT("fatal error"); + } + fprintf(stderr, "% 7.2f ", val); + } else { + fprintf(stderr, " "); + } + } + fprintf(stderr, "\n"); + } +} + +template +static void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, int split_k, int shader_size) { + VK_LOG_DEBUG("ggml_vk_test_matmul(" << m << ", " << n << ", " << k << ", " << batch << ", " << num_it << ", " << split_k << ", " << shader_size << ")"); + const size_t x_ne = m * k * batch; + const size_t y_ne = k * n * batch; + const size_t d_ne = m * n * batch; + + vk_pipeline p; + std::string shname; + if (shader_size == 0) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->a_s; + shname = "F32_ALIGNED_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->a_s; + shname = "F32_F16_ALIGNED_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_s; + shname = "F16_F32_ALIGNED_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->a_s; + shname = "F16_ALIGNED_S"; + } else { + GGML_ABORT("fatal error"); + } + } else if (shader_size == 1) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->a_m; + shname = "F32_ALIGNED_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->a_m; + shname = "F32_F16_ALIGNED_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_m; + shname = "F16_F32_ALIGNED_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->a_m; + shname = "F16_ALIGNED_M"; + } else { + GGML_ABORT("fatal error"); + } + } else if (shader_size == 2) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->a_l; + shname = "F32_ALIGNED_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->a_l; + shname = "F32_F16_ALIGNED_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_l; + shname = "F16_F32_ALIGNED_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->a_l; + shname = "F16_ALIGNED_L"; + } else { + GGML_ABORT("fatal error"); + } + } else { + GGML_ASSERT(0); + } + + const size_t kpad = ggml_vk_align_size(k, p->align); + + if (k != kpad) { + if (shader_size == 0) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->s; + shname = "F32_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->s; + shname = "F32_F16_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->s; + shname = "F16_F32_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->s; + shname = "F16_S"; + } + } else if (shader_size == 1) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->m; + shname = "F32_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->m; + shname = "F32_F16_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->m; + shname = "F16_F32_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->m; + shname = "F16_M"; + } + } else if (shader_size == 2) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->l; + shname = "F32_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->l; + shname = "F32_F16_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->l; + shname = "F16_F32_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->l; + shname = "F16_L"; + } + } + } + + ggml_pipeline_request_descriptor_sets(ctx->device, p, num_it); + if (split_k > 1) { + ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, num_it); + + if (ctx->prealloc_split_k == nullptr || ctx->prealloc_split_k->size < sizeof(float) * d_ne * split_k) { + // Resize buffer + if (ctx->prealloc_split_k != nullptr) { + ggml_vk_destroy_buffer(ctx->prealloc_split_k); + } + ctx->prealloc_split_k = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne * split_k, vk::MemoryPropertyFlagBits::eDeviceLocal); + } + } + + ggml_pipeline_allocate_descriptor_sets(ctx->device); + + vk_buffer d_X = ggml_vk_create_buffer_check(ctx->device, sizeof(X_TYPE) * x_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); + vk_buffer d_Y = ggml_vk_create_buffer_check(ctx->device, sizeof(Y_TYPE) * y_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); + vk_buffer d_D = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); + + X_TYPE* x = (X_TYPE *) malloc(sizeof(X_TYPE) * x_ne); + Y_TYPE* y = (Y_TYPE *) malloc(sizeof(Y_TYPE) * y_ne); + float* d = (float *) malloc(sizeof(float) * d_ne); + + for (size_t i = 0; i < x_ne; i++) { + if (std::is_same()) { + x[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; + // x[i] = 1.0f; + // x[i] = i + 1; + // x[i] = (i % k == i / k) ? 1.0f : 0.0f; + } else if (std::is_same()) { + x[i] = ggml_fp32_to_fp16((rand() / (float)RAND_MAX) * 2.0f - 1.0f); + // x[i] = ggml_fp32_to_fp16(1.0f); + // x[i] = ggml_fp32_to_fp16(i + 1); + // x[i] = ggml_fp32_to_fp16((i % k == i / k) ? 1.0f : 0.0f); + } else { + GGML_ABORT("fatal error"); + } + } + for (size_t i = 0; i < y_ne; i++) { + if (std::is_same()) { + y[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; + // y[i] = (i % k == i / k) ? 1.0f : 0.0f; + // y[i] = i + 1; + } else if (std::is_same()) { + y[i] = ggml_fp32_to_fp16((rand() / (float)RAND_MAX) * 2.0f - 1.0f); + // y[i] = ggml_fp32_to_fp16((i % k == i / k) ? 1.0f : 0.0f); + // y[i] = ggml_fp32_to_fp16(i + 1); + } else { + GGML_ABORT("fatal error"); + } + } + + ggml_vk_buffer_write(d_X, 0, x, sizeof(X_TYPE) * k * m * batch); + ggml_vk_buffer_write(d_Y, 0, y, sizeof(Y_TYPE) * k * n * batch); + + vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); + ggml_vk_ctx_begin(ctx->device, subctx); + for (size_t i = 0; i < num_it; i++) { + ggml_vk_matmul( + ctx, subctx, p, ggml_vk_subbuffer(d_X), ggml_vk_subbuffer(d_Y), ggml_vk_subbuffer(d_D), ggml_vk_subbuffer(ctx->prealloc_split_k), + m, n, k, + k, k, m, k*m, k*n, m*n, + split_k, batch, batch, batch, 1, 1 + ); + } + ggml_vk_ctx_end(subctx); + + auto begin = std::chrono::high_resolution_clock::now(); + ggml_vk_submit(subctx, ctx->fence); + VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_matmul waitForFences"); + ctx->device->device.resetFences({ ctx->fence }); + + auto end = std::chrono::high_resolution_clock::now(); + double time = std::chrono::duration_cast(end-begin).count() / 1000.0; + + // copy dst to host + ggml_vk_buffer_read(d_D, 0, d, sizeof(float) * d_ne); + + float * d_chk = (float *) malloc(sizeof(float) * d_ne); + + ggml_init_params iparams = { + /*.mem_size =*/ 1024*1024*1024, + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + + ggml_context * ggml_ctx = ggml_init(iparams); + + ggml_type src0_type; + ggml_type src1_type; + + if (std::is_same()) { + src0_type = GGML_TYPE_F32; + } else if (std::is_same()) { + src0_type = GGML_TYPE_F16; + } else { + GGML_ABORT("fatal error"); + } + if (std::is_same()) { + src1_type = GGML_TYPE_F32; + } else if (std::is_same()) { + src1_type = GGML_TYPE_F16; + } else { + GGML_ABORT("fatal error"); + } + + ggml_tensor * src0_ggml = ggml_new_tensor_3d(ggml_ctx, src0_type, k, m, batch); + ggml_tensor * src1_ggml = ggml_new_tensor_3d(ggml_ctx, src1_type, k, n, batch); + ggml_tensor * tensor_ggml = ggml_mul_mat(ggml_ctx, src0_ggml, src1_ggml); + + src0_ggml->data = x; + src1_ggml->data = y; + tensor_ggml->data = d_chk; + + ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); + ggml_build_forward_expand(cgraph, tensor_ggml); + + ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 1); + + ggml_free(ggml_ctx); + + double avg_err = 0.0; + int first_err_n = -1; + int first_err_m = -1; + int first_err_b = -1; + + for (size_t i = 0; i < m*n*batch; i++) { + double err = std::fabs(d[i] - d_chk[i]); + avg_err += err; + + if ((err > 0.05f || std::isnan(err)) && first_err_n == -1) { + first_err_b = i / (m * n); + first_err_n = (i % (m * n)) / m; + first_err_m = (i % (m * n)) % m; + } + } + + avg_err /= m * n; + + double tflops = 2.0*m*n*k*batch*num_it / (time / 1000.0) / (1000.0*1000.0*1000.0*1000.0); + + std::cerr << "TEST " << shname << " m=" << m << " n=" << n << " k=" << k << " batch=" << batch << " split_k=" << split_k << " matmul " << time / num_it << "ms " << tflops << " TFLOPS avg_err=" << avg_err << std::endl; + + if (avg_err > 0.1 || std::isnan(avg_err)) { + std::cerr << "m = " << first_err_m << " n = " << first_err_n << " b = " << first_err_b << std::endl; + std::cerr << "Actual result: " << std::endl << std::endl; + ggml_vk_print_matrix_area(d, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + std::cerr << "Expected result: " << std::endl << std::endl; + ggml_vk_print_matrix_area(d_chk, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + if (split_k > 1) { + float * split_k_buf = (float *) malloc(sizeof(float) * d_ne * split_k); + ggml_vk_buffer_read(ctx->prealloc_split_k, 0, split_k_buf, sizeof(float) * d_ne * split_k); + + std::cerr << "d_buf0: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf1: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf2: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + 2 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf3: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + 3 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + free(split_k_buf); + } + } + + free(d_chk); + + ggml_vk_queue_cleanup(ctx->device, ctx->device->transfer_queue); + ggml_vk_queue_cleanup(ctx->device, ctx->device->compute_queue); + + ggml_vk_destroy_buffer(d_X); + ggml_vk_destroy_buffer(d_Y); + ggml_vk_destroy_buffer(d_D); + + ggml_pipeline_cleanup(p); + ggml_pipeline_cleanup(ctx->device->pipeline_matmul_split_k_reduce); + + free(x); + free(y); + free(d); +} + +static void ggml_vk_print_tensor_area(const ggml_tensor * tensor, int i0, int i1, int i2, int i3) { + if (tensor->type != GGML_TYPE_F32 && tensor->type != GGML_TYPE_F16) { + return; + } + i0 = std::max(i0, 5); + i1 = std::max(i1, 5); + i2 = std::max(i2, 0); + i3 = std::max(i3, 0); + fprintf(stderr, " "); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + fprintf(stderr, "%7d ", idx1); + } + fprintf(stderr, "\n"); + for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { + fprintf(stderr, "%7d: ", idx0); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + if (idx0 >= 0 && idx0 < tensor->ne[0] && idx1 >= 0 && idx1 < tensor->ne[1] && i2 >= 0 && i2 < tensor->ne[2] && i3 >= 0 && i3 < tensor->ne[3]) { + float val; + if (tensor->type == GGML_TYPE_F32) { + val = *(float *) ((char *) tensor->data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); + } else if (tensor->type == GGML_TYPE_F16) { + val = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) tensor->data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0])); + } else { + GGML_ABORT("fatal error"); + } + fprintf(stderr, "% 7.2f ", val); + } else { + fprintf(stderr, " "); + } + } + fprintf(stderr, "\n"); + } +} + +static void ggml_vk_quantize_data(const float * from, void * to, size_t ne, ggml_type quant) { + ggml_quantize_chunk(quant, from, to, 0, 1, ne, nullptr); +} + +static void ggml_vk_dequantize_data(const void * from, float * to, size_t ne, ggml_type quant) { + if (quant == GGML_TYPE_F32) { + memcpy(to, from, sizeof(float) * ne); + return; + } + + const auto * tt = ggml_get_type_traits(quant); + + ggml_to_float_t dequant_fn = tt->to_float; + + dequant_fn(from, to, ne); +} + +static void ggml_vk_test_dequant(ggml_backend_vk_context * ctx, size_t ne, ggml_type quant) { + VK_LOG_DEBUG("ggml_vk_test_dequant(" << ne << ")"); + const size_t x_sz = sizeof(float) * ne; + const size_t x_sz_f16 = sizeof(ggml_fp16_t) * ne; + const size_t qx_sz = ne * ggml_type_size(quant)/ggml_blck_size(quant); + float * x = (float *) malloc(x_sz); + void * qx = malloc(qx_sz); + vk_buffer qx_buf = ggml_vk_create_buffer_check(ctx->device, qx_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); + vk_buffer x_buf = ggml_vk_create_buffer_check(ctx->device, x_sz_f16, vk::MemoryPropertyFlagBits::eDeviceLocal); + float * x_ref = (float *) malloc(x_sz); + ggml_fp16_t * x_chk = (ggml_fp16_t *) malloc(x_sz_f16); + + for (size_t i = 0; i < ne; i++) { + x[i] = rand() / (float)RAND_MAX; + } + + vk_pipeline p = ggml_vk_get_to_fp16(ctx, quant); + + ggml_vk_quantize_data(x, qx, ne, quant); + ggml_vk_dequantize_data(qx, x_ref, ne, quant); + + ggml_pipeline_request_descriptor_sets(ctx->device, p, 1); + + ggml_pipeline_allocate_descriptor_sets(ctx->device); + + ggml_vk_buffer_write(qx_buf, 0, qx, qx_sz); + + vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); + ggml_vk_ctx_begin(ctx->device, subctx); + const std::vector pc = { 1, (uint32_t)ne, (uint32_t)ne, (uint32_t)ne, (uint32_t)ne }; + ggml_vk_dispatch_pipeline(ctx, subctx, p, { vk_subbuffer{ qx_buf, 0, qx_sz }, vk_subbuffer{ x_buf, 0, x_sz_f16 } }, pc.size() * sizeof(int), pc.data(), { (uint32_t)ne, 1, 1}); + ggml_vk_ctx_end(subctx); + + auto begin = std::chrono::high_resolution_clock::now(); + + ggml_vk_submit(subctx, ctx->fence); + VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences"); + ctx->device->device.resetFences({ ctx->fence }); + + auto end = std::chrono::high_resolution_clock::now(); + + double ms_dequant = std::chrono::duration_cast(end-begin).count() / 1000.0; + ggml_vk_buffer_read(x_buf, 0, x_chk, x_sz_f16); + + int first_err = -1; + + double avg_err = 0.0; + for (size_t i = 0; i < ne; i++) { + double error = std::fabs(x_ref[i] - ggml_fp16_to_fp32(x_chk[i])); + avg_err += error; + + if (first_err < 0 && error > 0.05) { + first_err = i; + } + } + + avg_err /= ne; + + std::cerr << "TEST DEQUANT " << ggml_type_name(quant) << " time=" << ms_dequant << "ms avg_err=" << avg_err << std::endl; + + if (avg_err > 0.1) { + std::cerr << "first_error = " << first_err << std::endl; + std::cerr << "Actual result: " << std::endl << std::endl; + for (int i = std::max(0, first_err - 5); i < std::min((int)ne, first_err + 5); i++) { + std::cerr << ggml_fp16_to_fp32(x_chk[i]) << ", "; + } + std::cerr << std::endl << "Expected result: " << std::endl << std::endl; + for (int i = std::max(0, first_err - 5); i < std::min((int)ne, first_err + 5); i++) { + std::cerr << x_ref[i] << ", "; + } + std::cerr << std::endl; + } + + ggml_vk_destroy_buffer(x_buf); + ggml_vk_destroy_buffer(qx_buf); + + free(x); + free(qx); + free(x_ref); + free(x_chk); +} + +static void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, size_t split_k, size_t shader_size, ggml_type quant) { + VK_LOG_DEBUG("ggml_vk_test_dequant_matmul(" << m << ", " << n << ", " << k << ", " << batch << ", " << num_it << ", " << split_k << ", " << ggml_type_name(quant) << ")"); + const size_t x_ne = m * k * batch; + const size_t y_ne = k * n * batch; + const size_t d_ne = m * n * batch; + + vk_pipeline p; + std::string shname; + if (shader_size == 0) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_s : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_s; + shname = std::string(ggml_type_name(quant)) + "_ALIGNED_S"; + } else if (shader_size == 1) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_m : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_m; + shname = std::string(ggml_type_name(quant)) + "_ALIGNED_M"; + } else if (shader_size == 2) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_l : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_l; + shname = std::string(ggml_type_name(quant)) + "_ALIGNED_L"; + } else { + GGML_ASSERT(0); + } + + const size_t kpad = ggml_vk_align_size(k, p->align); + + if (k != kpad) { + if (shader_size == 0) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->s : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->s; + shname = std::string(ggml_type_name(quant)) + "_S"; + } else if (shader_size == 1) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->m : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->m; + shname = std::string(ggml_type_name(quant)) + "_M"; + } else if (shader_size == 2) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->l : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->l; + shname = std::string(ggml_type_name(quant)) + "_L"; + } else { + GGML_ASSERT(0); + } + } + + const size_t x_sz = sizeof(float) * x_ne; + const size_t y_sz = sizeof(float) * y_ne; + const size_t qx_sz = x_ne * ggml_type_size(quant)/ggml_blck_size(quant); + const size_t d_sz = sizeof(float) * d_ne; + float * x = (float *) malloc(x_sz); + float * y = (float *) malloc(y_sz); + void * qx = malloc(qx_sz); + vk_buffer qx_buf = ggml_vk_create_buffer_check(ctx->device, qx_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); + vk_buffer y_buf = ggml_vk_create_buffer_check(ctx->device, y_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); + vk_buffer d_buf = ggml_vk_create_buffer_check(ctx->device, d_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); + float * d = (float *) malloc(d_sz); + float * d_chk = (float *) malloc(d_sz); + + for (size_t i = 0; i < x_ne; i++) { + x[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; + } + + ggml_vk_quantize_data(x, qx, x_ne, quant); + + for (size_t i = 0; i < y_ne; i++) { + // y[i] = rand() / (float)RAND_MAX; + y[i] = (i % k == i / k) ? 1.0f : 0.0f; + } + + ggml_pipeline_request_descriptor_sets(ctx->device, p, num_it); + if (split_k > 1) { + ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, num_it); + + if (ctx->prealloc_split_k == nullptr || ctx->prealloc_split_k->size < sizeof(float) * d_ne * split_k) { + // Resize buffer + if (ctx->prealloc_split_k != nullptr) { + ggml_vk_destroy_buffer(ctx->prealloc_split_k); + } + ctx->prealloc_split_k = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne * split_k, vk::MemoryPropertyFlagBits::eDeviceLocal); + } + } + + ggml_pipeline_allocate_descriptor_sets(ctx->device); + + ggml_vk_buffer_write(qx_buf, 0, qx, qx_sz); + ggml_vk_buffer_write(y_buf, 0, y, y_sz); + + vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); + ggml_vk_ctx_begin(ctx->device, subctx); + for (size_t i = 0; i < num_it; i++) { + ggml_vk_matmul( + ctx, subctx, p, ggml_vk_subbuffer(qx_buf), ggml_vk_subbuffer(y_buf), ggml_vk_subbuffer(d_buf), ggml_vk_subbuffer(ctx->prealloc_split_k), + m, n, k, + k, k, m, k*m, k*n, m*n, + split_k, batch, batch, batch, 1, 1 + ); + } + ggml_vk_ctx_end(subctx); + + auto begin = std::chrono::high_resolution_clock::now(); + + ggml_vk_submit(subctx, ctx->fence); + VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences"); + ctx->device->device.resetFences({ ctx->fence }); + + auto end = std::chrono::high_resolution_clock::now(); + + double time_ms = std::chrono::duration_cast(end-begin).count() / 1000.0; + ggml_vk_buffer_read(d_buf, 0, d, d_sz); + + ggml_init_params iparams = { + /*.mem_size =*/ 1024*1024*1024, + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + + ggml_context * ggml_ctx = ggml_init(iparams); + + ggml_tensor * src0_ggml = ggml_new_tensor_3d(ggml_ctx, quant, k, m, batch); + ggml_tensor * src1_ggml = ggml_new_tensor_3d(ggml_ctx, GGML_TYPE_F32, k, n, batch); + ggml_tensor * tensor_ggml = ggml_mul_mat(ggml_ctx, src0_ggml, src1_ggml); + + src0_ggml->data = qx; + src1_ggml->data = y; + tensor_ggml->data = d_chk; + + ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); + ggml_build_forward_expand(cgraph, tensor_ggml); + + ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 1); + + ggml_free(ggml_ctx); + + double avg_err = 0.0; + int first_err_n = -1; + int first_err_m = -1; + int first_err_b = -1; + + for (size_t i = 0; i < m*n*batch; i++) { + double err = std::fabs(d[i] - d_chk[i]); + avg_err += err; + + if ((err > 0.05f || std::isnan(err)) && first_err_n == -1) { + first_err_b = i / (m * n); + first_err_n = (i % (m * n)) / m; + first_err_m = (i % (m * n)) % m; + } + } + + avg_err /= m * n; + + double tflops = 2.0*m*n*k*batch*num_it / (time_ms / 1000.0) / (1000.0*1000.0*1000.0*1000.0); + + std::cerr << "TEST MMQ " << shname << " m=" << m << " n=" << n << " k=" << k << " batch=" << batch << " split_k=" << split_k << " matmul " << time_ms / num_it << "ms " << tflops << " TFLOPS avg_err=" << avg_err << std::endl; + + if (avg_err > 0.01 || std::isnan(avg_err)) { + std::cerr << "m = " << first_err_m << " n = " << first_err_n << " b = " << first_err_b << std::endl; + std::cerr << "Actual result: " << std::endl << std::endl; + ggml_vk_print_matrix_area(d, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + std::cerr << std::endl; + std::cerr << "Expected result: " << std::endl << std::endl; + ggml_vk_print_matrix_area(d_chk, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + if (split_k > 1) { + float * split_k_buf = (float *) malloc(sizeof(float) * d_ne * split_k); + ggml_vk_buffer_read(ctx->prealloc_split_k, 0, split_k_buf, sizeof(float) * d_ne * split_k); + + std::cerr << "d_buf0: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf1: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf2: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + 2 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf3: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + 3 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + free(split_k_buf); + } + } + + ggml_vk_destroy_buffer(qx_buf); + ggml_vk_destroy_buffer(y_buf); + ggml_vk_destroy_buffer(d_buf); + + free(x); + free(qx); + free(y); + free(d); + free(d_chk); +} +#endif + +static void ggml_vk_preallocate_buffers(ggml_backend_vk_context * ctx) { +#if defined(GGML_VULKAN_RUN_TESTS) + const std::vector vals { + 512, 512, 128, + 128, 512, 512, + 4096, 512, 4096, + 11008, 512, 4096, + 4096, 512, 11008, + 32000, 512, 4096, + 8, 8, 8, + 100, 46, 576, + 623, 111, 128, + 100, 46, 558, + 512, 1, 256, + 128, 110, 622, + 511, 511, 127, + 511, 511, 7, + 511, 511, 17, + 49, 49, 128, + 128, 49, 49, + 4096, 49, 4096, + }; + const size_t num_it = 100; + + for (size_t i = 0; i < vals.size(); i += 3) { + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2); + std::cerr << '\n'; + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2); + std::cerr << '\n'; + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2); + std::cerr << '\n' << std::endl; + + if (vals[i + 2] % 32 == 0) { + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2, GGML_TYPE_Q4_0); + std::cerr << '\n'; + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2, GGML_TYPE_Q4_0); + std::cerr << '\n'; + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2, GGML_TYPE_Q4_0); + std::cerr << '\n' << std::endl; + } + + if (vals[i + 2] % 256 == 0) { + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2, GGML_TYPE_Q4_K); + std::cerr << '\n'; + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2, GGML_TYPE_Q4_K); + std::cerr << '\n'; + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2, GGML_TYPE_Q4_K); + std::cerr << '\n' << std::endl; + } + } + + GGML_ABORT("fatal error"); +#endif + + if (ctx->prealloc_x == nullptr || (ctx->prealloc_size_x > 0 && ctx->prealloc_x->size < ctx->prealloc_size_x)) { + VK_LOG_MEMORY("ggml_vk_preallocate_buffers(x_size: " << ctx->prealloc_size_x << ")"); + // Resize buffer + if (ctx->prealloc_x != nullptr) { + ggml_vk_destroy_buffer(ctx->prealloc_x); + } + ctx->prealloc_x = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_x); + } + if (ctx->prealloc_y == nullptr || (ctx->prealloc_size_y > 0 && ctx->prealloc_y->size < ctx->prealloc_size_y)) { + VK_LOG_MEMORY("ggml_vk_preallocate_buffers(y_size: " << ctx->prealloc_size_y << ")"); + // Resize buffer + if (ctx->prealloc_y != nullptr) { + ggml_vk_destroy_buffer(ctx->prealloc_y); + } + ctx->prealloc_y = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_y); + } + if (ctx->prealloc_split_k == nullptr || (ctx->prealloc_size_split_k > 0 && ctx->prealloc_split_k->size < ctx->prealloc_size_split_k)) { + VK_LOG_MEMORY("ggml_vk_preallocate_buffers(split_k_size: " << ctx->prealloc_size_split_k << ")"); + // Resize buffer + if (ctx->prealloc_split_k != nullptr) { + ggml_vk_destroy_buffer(ctx->prealloc_split_k); + } + ctx->prealloc_split_k = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_split_k); + } +} + +static bool ggml_vk_compute_forward(ggml_backend_vk_context* ctx, ggml_tensor* tensor, int tensor_idx, bool use_fence); + +// Returns true if node has enqueued work into the queue, false otherwise +// If submit is true the current all operations queued so far are being submitted to Vulkan to overlap cmdlist creation and GPU execution. +static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_tensor * node, int node_idx, ggml_tensor *node_begin, int node_idx_begin, bool dryrun, bool last_node, bool submit){ + if (ggml_is_empty(node) || !node->buffer) { + return false; + } + + VK_LOG_DEBUG("ggml_vk_build_graph(" << node << ", " << ggml_op_name(node->op) << ")"); + ctx->semaphore_idx = 0; + + const ggml_tensor * src0 = node->src[0]; + const ggml_tensor * src1 = node->src[1]; + const ggml_tensor * src2 = node->src[2]; + const ggml_tensor * src3 = node->src[3]; + + switch (node->op) { + // Return on empty ops to avoid generating a compute_ctx and setting exit_tensor + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_NONE: + return false; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(node)) { + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_GELU: + case GGML_UNARY_OP_GELU_QUICK: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_TANH: + break; + default: + return false; + } + break; + case GGML_OP_REPEAT: + case GGML_OP_GET_ROWS: + case GGML_OP_ADD: + case GGML_OP_ACC: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_DUP: + case GGML_OP_NORM: + case GGML_OP_GROUP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_SOFT_MAX: + case GGML_OP_ROPE: + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + case GGML_OP_ARGSORT: + case GGML_OP_SUM_ROWS: + case GGML_OP_IM2COL: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_POOL_2D: + case GGML_OP_RWKV_WKV6: + case GGML_OP_LEAKY_RELU: + case GGML_OP_FLASH_ATTN_EXT: + break; + default: + std::cerr << "ggml_vulkan: Error: Missing op: " << ggml_op_name(node->op) << std::endl; + GGML_ABORT("fatal error"); + return false; + } + + vk_context compute_ctx; + + if (!dryrun) { + if (ctx->compute_ctx.expired()) { + compute_ctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); + ctx->compute_ctx = compute_ctx; + ggml_vk_ctx_begin(ctx->device, compute_ctx); + } else { + compute_ctx = ctx->compute_ctx.lock(); + } + } else { + switch (node->op) { + case GGML_OP_REPEAT: + case GGML_OP_ACC: + case GGML_OP_GET_ROWS: + case GGML_OP_ADD: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_DUP: + case GGML_OP_NORM: + case GGML_OP_GROUP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_UNARY: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_SOFT_MAX: + case GGML_OP_ROPE: + case GGML_OP_ARGSORT: + case GGML_OP_SUM_ROWS: + case GGML_OP_IM2COL: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_POOL_2D: + case GGML_OP_LEAKY_RELU: + { + // These operations all go through ggml_vk_op_f32, so short-circuit and + // do the only thing needed for the dryrun. + vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, src2, node, node->op); + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + return false; + } + default: + break; + } + } + + switch (node->op) { + case GGML_OP_REPEAT: + ggml_vk_repeat(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_ACC: + ggml_vk_acc(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_GET_ROWS: + ggml_vk_get_rows(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_ADD: + ggml_vk_add(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_MUL: + ggml_vk_mul(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_DIV: + ggml_vk_div(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_CONCAT: + ggml_vk_concat(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_UPSCALE: + ggml_vk_upscale(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_SCALE: + ggml_vk_scale(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_SQR: + ggml_vk_sqr(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_SIN: + ggml_vk_sin(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_COS: + ggml_vk_cos(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_CLAMP: + ggml_vk_clamp(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_PAD: + ggml_vk_pad(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_DUP: + ggml_vk_cpy(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_NORM: + ggml_vk_norm(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_GROUP_NORM: + ggml_vk_group_norm(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_RMS_NORM: + ggml_vk_rms_norm(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(node)) { + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_GELU: + case GGML_UNARY_OP_GELU_QUICK: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_TANH: + ggml_vk_unary(ctx, compute_ctx, src0, node, dryrun); + break; + default: + return false; + } + break; + case GGML_OP_DIAG_MASK_INF: + ggml_vk_diag_mask_inf(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_SOFT_MAX: + ggml_vk_soft_max(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_ROPE: + ggml_vk_rope(ctx, compute_ctx, src0, src1, src2, node, dryrun); + + break; + case GGML_OP_ARGSORT: + ggml_vk_argsort(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_SUM_ROWS: + ggml_vk_sum_rows(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_IM2COL: + ggml_vk_im2col(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_TIMESTEP_EMBEDDING: + ggml_vk_timestep_embedding(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_POOL_2D: + ggml_vk_pool_2d(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_LEAKY_RELU: + ggml_vk_leaky_relu(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_MUL_MAT: + ggml_vk_mul_mat(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_MUL_MAT_ID: + ggml_vk_mul_mat_id(ctx, compute_ctx, src0, src1, src2, node, dryrun); + + break; + + case GGML_OP_FLASH_ATTN_EXT: + ggml_vk_flash_attn(ctx, compute_ctx, src0, src1, src2, src3, node, dryrun); + + break; + + case GGML_OP_RWKV_WKV6: + ggml_vk_rwkv_wkv6(ctx, compute_ctx, node, dryrun); + + break; + default: + return false; + } + + if (dryrun) { + return false; + } + + ctx->tensor_ctxs[node_idx] = compute_ctx; + +#if defined(GGML_VULKAN_CHECK_RESULTS) || defined(GGML_VULKAN_PERF) + // Force context reset on each node so that each tensor ends up in its own context + // and can be run and compared to its CPU equivalent separately + last_node = true; +#endif + + if (submit || last_node) { + ggml_vk_ctx_end(compute_ctx); + + // TODO probably it'd be better to pass a exit_node flag to ggml_vk_compute_forward + if (last_node) { + compute_ctx->exit_tensor_idx = node_idx_begin; + } + else { + compute_ctx->exit_tensor_idx = -1; + } + + ctx->compute_ctx.reset(); + + bool ok = ggml_vk_compute_forward(ctx, node_begin, node_idx_begin, false); + if (!ok) { + if (node->op == GGML_OP_UNARY) { + std::cerr << __func__ << ": error: op not supported UNARY " << node->name << " (" << ggml_unary_op_name(static_cast(node->op_params[0])) << ")" << std::endl; + } + else { + std::cerr << __func__ << ": error: op not supported " << node->name << " (" << ggml_op_name(node->op) << ")" << std::endl; + } + } + + } + return true; +} + +static bool ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_tensor * tensor, int tensor_idx, bool use_fence = true){ + ggml_backend_buffer * buf = nullptr; + + switch (tensor->op) { + case GGML_OP_ADD: + case GGML_OP_ACC: + case GGML_OP_GET_ROWS: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_DUP: + case GGML_OP_NORM: + case GGML_OP_GROUP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_SOFT_MAX: + case GGML_OP_ROPE: + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_NONE: + case GGML_OP_ARGSORT: + case GGML_OP_SUM_ROWS: + case GGML_OP_IM2COL: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_POOL_2D: + case GGML_OP_RWKV_WKV6: + case GGML_OP_LEAKY_RELU: + case GGML_OP_REPEAT: + buf = tensor->buffer; + + break; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(tensor)) { + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_GELU: + case GGML_UNARY_OP_GELU_QUICK: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_TANH: + buf = tensor->buffer; + break; + default: + return false; + } + break; + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + case GGML_OP_FLASH_ATTN_EXT: + buf = tensor->buffer; + + break; + default: + return false; + } + + if (buf == nullptr) { + return false; + } + + VK_LOG_DEBUG("ggml_vk_compute_forward(" << tensor << ", name=" << tensor->name << ", op=" << ggml_op_name(tensor->op) << ", type=" << tensor->type << ", ne0=" << tensor->ne[0] << ", ne1=" << tensor->ne[1] << ", ne2=" << tensor->ne[2] << ", ne3=" << tensor->ne[3] << ", nb0=" << tensor->nb[0] << ", nb1=" << tensor->nb[1] << ", nb2=" << tensor->nb[2] << ", nb3=" << tensor->nb[3] << ", view_src=" << tensor->view_src << ", view_offs=" << tensor->view_offs << ")"); + + vk_context subctx = ctx->tensor_ctxs[tensor_idx].lock(); + + // always wait for the GPU work to be done for the last submit + if (tensor_idx == subctx->exit_tensor_idx) { + use_fence = true; + } + + // Only run if ctx hasn't been submitted yet + if (!subctx->seqs.empty()) { +#ifdef GGML_VULKAN_CHECK_RESULTS + ggml_vk_check_results_0(tensor); + use_fence = true; +#endif + + // Do staging buffer copies + for (auto& cpy : subctx->in_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + + ggml_vk_submit(subctx, use_fence ? ctx->fence : vk::Fence{}); + + if (use_fence) { + VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_compute_forward waitForFences"); + + ctx->device->device.resetFences({ ctx->fence }); + } +#ifdef GGML_VULKAN_CHECK_RESULTS + ggml_vk_check_results_1(tensor); +#endif + } + + if (tensor_idx == subctx->exit_tensor_idx) { + // Do staging buffer copies + for (auto& cpy : subctx->out_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + subctx->in_memcpys.clear(); + subctx->out_memcpys.clear(); + } + + return true; +} + +// Clean up after graph processing is done +static void ggml_vk_graph_cleanup(ggml_backend_vk_context * ctx) { + VK_LOG_DEBUG("ggml_vk_graph_cleanup()"); + for (auto& buffer : ctx->gc.temp_buffers) { + ggml_vk_pool_free(ctx, buffer); + } + ctx->gc.temp_buffers.clear(); + + for (auto& dsr : ctx->device->pipeline_descriptor_set_requirements) { + vk_pipeline_ref plr = ctx->device->pipelines[dsr.first]; + + if (plr.expired()) { + continue; + } + + vk_pipeline pl = plr.lock(); + ggml_pipeline_cleanup(pl); + } + + ggml_vk_queue_cleanup(ctx->device, ctx->device->compute_queue); + ggml_vk_queue_cleanup(ctx->device, ctx->device->transfer_queue); + + for (size_t i = 0; i < ctx->gc.semaphores.size(); i++) { + ctx->device->device.destroySemaphore({ ctx->gc.semaphores[i].s }); + } + ctx->gc.semaphores.clear(); + + for (size_t i = 0; i < ctx->gc.tl_semaphores.size(); i++) { + ctx->device->device.destroySemaphore({ ctx->gc.tl_semaphores[i].s }); + } + ctx->gc.tl_semaphores.clear(); + ctx->semaphore_idx = 0; + + ctx->event_idx = 0; + + for (auto& event : ctx->gc.events) { + ctx->device->device.resetEvent(event); + } + + ctx->tensor_ctxs.clear(); + ctx->gc.contexts.clear(); + ctx->device->pipeline_descriptor_set_requirements.clear(); +} + +// Clean up on backend free +static void ggml_vk_cleanup(ggml_backend_vk_context * ctx) { + VK_LOG_DEBUG("ggml_vk_cleanup(" << ctx->name << ")"); + ggml_vk_graph_cleanup(ctx); + + ggml_vk_destroy_buffer(ctx->prealloc_x); + ggml_vk_destroy_buffer(ctx->prealloc_y); + ggml_vk_destroy_buffer(ctx->prealloc_split_k); + + for (auto& buffer : ctx->buffer_pool) { + ggml_vk_destroy_buffer(buffer); + } + + ctx->prealloc_size_x = 0; + ctx->prealloc_size_y = 0; + ctx->prealloc_size_split_k = 0; + + for (auto& event : ctx->gc.events) { + ctx->device->device.destroyEvent(event); + } + ctx->gc.events.clear(); + + ctx->device->device.destroyFence(ctx->fence); +} + +static int ggml_vk_get_device_count() { + ggml_vk_instance_init(); + + return vk_instance.device_indices.size(); +} + +static void ggml_vk_get_device_description(int device, char * description, size_t description_size) { + ggml_vk_instance_init(); + + std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); + + vk::PhysicalDeviceProperties props; + devices[device].getProperties(&props); + + snprintf(description, description_size, "%s", props.deviceName.data()); +} + +// backend interface + +#define UNUSED GGML_UNUSED + +// device backend + +static bool ggml_backend_buffer_is_vk(ggml_backend_buffer_t buffer) { + return buffer->buft->iface.get_name == ggml_backend_vk_buffer_type_name; +} + +static void ggml_backend_vk_buffer_free_buffer(ggml_backend_buffer_t buffer) { + VK_LOG_MEMORY("ggml_backend_vk_buffer_free_buffer()"); + ggml_backend_vk_buffer_context * ctx = (ggml_backend_vk_buffer_context *)buffer->context; + ggml_vk_destroy_buffer(ctx->dev_buffer); + delete ctx; +} + +static void * ggml_backend_vk_buffer_get_base(ggml_backend_buffer_t buffer) { + return vk_ptr_base; + + UNUSED(buffer); +} + +static void ggml_backend_vk_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + VK_LOG_DEBUG("ggml_backend_vk_buffer_init_tensor(" << buffer << " (" << buffer->context << "), " << tensor << ")"); + if (tensor->view_src != nullptr) { + GGML_ASSERT(tensor->view_src->buffer->buft == buffer->buft); + } +} + +static void ggml_backend_vk_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + VK_LOG_DEBUG("ggml_backend_vk_buffer_set_tensor(" << buffer << ", " << tensor << ", " << data << ", " << offset << ", " << size << ")"); + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)buffer->context; + vk_buffer buf = buf_ctx->dev_buffer; + + ggml_vk_buffer_write(buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); +} + +static void ggml_backend_vk_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + VK_LOG_DEBUG("ggml_backend_vk_buffer_get_tensor(" << buffer << ", " << tensor << ", " << data << ", " << offset << ", " << size << ")"); + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)buffer->context; + + vk_buffer buf = buf_ctx->dev_buffer; + + ggml_vk_buffer_read(buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); +} + +static bool ggml_backend_vk_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) { + if (ggml_backend_buffer_is_vk(src->buffer)) { + ggml_backend_vk_buffer_context * src_buf_ctx = (ggml_backend_vk_buffer_context *)src->buffer->context; + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + + vk_buffer src_buf = src_buf_ctx->dev_buffer; + vk_buffer dst_buf = dst_buf_ctx->dev_buffer; + + ggml_vk_buffer_copy(dst_buf, vk_tensor_offset(dst) + dst->view_offs, src_buf, vk_tensor_offset(src) + src->view_offs, ggml_nbytes(src)); + + return true; + } + return false; + + UNUSED(buffer); +} + +static void ggml_backend_vk_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + ggml_backend_vk_buffer_context * ctx = (ggml_backend_vk_buffer_context *)buffer->context; + + ggml_vk_buffer_memset(ctx->dev_buffer, 0, value, buffer->size); +} + +static ggml_backend_buffer_i ggml_backend_vk_buffer_interface = { + /* .free_buffer = */ ggml_backend_vk_buffer_free_buffer, + /* .get_base = */ ggml_backend_vk_buffer_get_base, + /* .init_tensor = */ ggml_backend_vk_buffer_init_tensor, + /* .memset_tensor = */ NULL, + /* .set_tensor = */ ggml_backend_vk_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_vk_buffer_get_tensor, + /* .cpy_tensor = */ ggml_backend_vk_buffer_cpy_tensor, + /* .clear = */ ggml_backend_vk_buffer_clear, + /* .reset = */ NULL, +}; + +// vk buffer type +static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft) { + ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *)buft->context; + + return ctx->name.c_str(); +} + +static ggml_backend_buffer_t ggml_backend_vk_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + VK_LOG_MEMORY("ggml_backend_vk_buffer_type_alloc_buffer(" << size << ")"); + ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; + + vk_buffer dev_buffer = nullptr; + try { + dev_buffer = ggml_vk_create_buffer_device(ctx->device, size); + } catch (const vk::SystemError& e) { + return nullptr; + } + + ggml_backend_vk_buffer_context * bufctx = new ggml_backend_vk_buffer_context(ctx->device, std::move(dev_buffer), ctx->name); + + return ggml_backend_buffer_init(buft, ggml_backend_vk_buffer_interface, bufctx, size); +} + +static size_t ggml_backend_vk_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; + return ctx->device->properties.limits.minStorageBufferOffsetAlignment; +} + +static size_t ggml_backend_vk_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { + ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; + return ctx->device->max_memory_allocation_size; +} + +static size_t ggml_backend_vk_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + return ggml_nbytes(tensor); + + UNUSED(buft); +} + +ggml_backend_buffer_type_t ggml_backend_vk_buffer_type(size_t dev_num) { + ggml_vk_instance_init(); + + VK_LOG_DEBUG("ggml_backend_vk_buffer_type(" << dev_num << ")"); + + vk_device dev = ggml_vk_get_device(dev_num); + + return &dev->buffer_type; +} + +// host buffer type + +static const char * ggml_backend_vk_host_buffer_type_name(ggml_backend_buffer_type_t buft) { + return GGML_VK_NAME "_Host"; + + UNUSED(buft); +} + +static const char * ggml_backend_vk_host_buffer_name(ggml_backend_buffer_t buffer) { + return GGML_VK_NAME "_Host"; + + UNUSED(buffer); +} + +static void ggml_backend_vk_host_buffer_free_buffer(ggml_backend_buffer_t buffer) { + VK_LOG_MEMORY("ggml_backend_vk_host_buffer_free_buffer()"); + ggml_vk_host_free(vk_instance.devices[0], buffer->context); +} + +static ggml_backend_buffer_t ggml_backend_vk_host_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + VK_LOG_MEMORY("ggml_backend_vk_host_buffer_type_alloc_buffer(" << size << ")"); + + size += 32; // Behave like the CPU buffer type + void * ptr = nullptr; + try { + ptr = ggml_vk_host_malloc(vk_instance.devices[0], size); + } catch (vk::SystemError& e) { + std::cerr << "ggml_vulkan: Failed to allocate pinned memory." << std::endl; + std::cerr << "ggml_vulkan: " << e.what() << std::endl; + // fallback to cpu buffer + return ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); + } + + ggml_backend_buffer_t buffer = ggml_backend_cpu_buffer_from_ptr(ptr, size); + buffer->buft = buft; + buffer->iface.free_buffer = ggml_backend_vk_host_buffer_free_buffer; + + return buffer; + + UNUSED(buft); +} + +static size_t ggml_backend_vk_host_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + return vk_instance.devices[0]->properties.limits.minMemoryMapAlignment; + + UNUSED(buft); +} + +// Should be changed to return device-specific host buffer type +// but that probably requires changes in llama.cpp +ggml_backend_buffer_type_t ggml_backend_vk_host_buffer_type() { + static struct ggml_backend_buffer_type ggml_backend_vk_buffer_type_host = { + /* .iface = */ { + /* .get_name = */ ggml_backend_vk_host_buffer_type_name, + /* .alloc_buffer = */ ggml_backend_vk_host_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_vk_host_buffer_type_get_alignment, + /* .get_max_size = */ NULL, // defaults to SIZE_MAX + /* .get_alloc_size = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size, + /* .is_host = */ ggml_backend_cpu_buffer_type()->iface.is_host, + }, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), 0), + /* .context = */ nullptr, + }; + + // Make sure device 0 is initialized + ggml_vk_instance_init(); + ggml_vk_get_device(0); + + return &ggml_backend_vk_buffer_type_host; +} + + +// backend + +static const char * ggml_backend_vk_name(ggml_backend_t backend) { + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + + return ctx->name.c_str(); +} + +static void ggml_backend_vk_free(ggml_backend_t backend) { + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + VK_LOG_DEBUG("ggml_backend_vk_free(" << ctx->name << ")"); + + ggml_vk_cleanup(ctx); + + delete ctx; + delete backend; +} + +static ggml_backend_buffer_type_t ggml_backend_vk_get_default_buffer_type(ggml_backend_t backend) { + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + + return &ctx->device->buffer_type; +} + +static void ggml_backend_vk_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + VK_LOG_DEBUG("ggml_backend_vk_set_tensor_async(" << size << ")"); + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + GGML_ASSERT((tensor->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || tensor->buffer->buft == ggml_backend_vk_host_buffer_type()) && "unsupported buffer type"); + + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; + + vk_context transfer_ctx; + + if (ctx->transfer_ctx.expired()) { + // Initialize new transfer context + transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); + ctx->transfer_ctx = transfer_ctx; + ggml_vk_ctx_begin(ctx->device, transfer_ctx); + } else { + transfer_ctx = ctx->transfer_ctx.lock(); + } + + vk_buffer buf = buf_ctx->dev_buffer; + + ggml_vk_buffer_write_async(transfer_ctx, buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); +} + +static void ggml_backend_vk_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + VK_LOG_DEBUG("ggml_backend_vk_get_tensor_async(" << size << ")"); + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + GGML_ASSERT((tensor->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || tensor->buffer->buft == ggml_backend_vk_host_buffer_type()) && "unsupported buffer type"); + + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; + + vk_context transfer_ctx; + + if (ctx->transfer_ctx.expired()) { + // Initialize new transfer context + transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); + ctx->transfer_ctx = transfer_ctx; + ggml_vk_ctx_begin(ctx->device, transfer_ctx); + } else { + transfer_ctx = ctx->transfer_ctx.lock(); + } + + vk_buffer buf = buf_ctx->dev_buffer; + + ggml_vk_buffer_read_async(transfer_ctx, buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); +} + +static bool ggml_backend_vk_cpy_tensor_async(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst) { + VK_LOG_DEBUG("ggml_backend_vk_cpy_tensor_async()"); + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + if ((dst->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || dst->buffer->buft == ggml_backend_vk_host_buffer_type()) && ggml_backend_buffer_is_vk(src->buffer)) { + ggml_backend_vk_buffer_context * src_buf_ctx = (ggml_backend_vk_buffer_context *)src->buffer->context; + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + + vk_context transfer_ctx; + + if (ctx->transfer_ctx.expired()) { + // Initialize new transfer context + transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); + ctx->transfer_ctx = transfer_ctx; + ggml_vk_ctx_begin(ctx->device, transfer_ctx); + } else { + transfer_ctx = ctx->transfer_ctx.lock(); + } + + vk_buffer src_buf = src_buf_ctx->dev_buffer; + vk_buffer dst_buf = dst_buf_ctx->dev_buffer; + + ggml_vk_buffer_copy_async(transfer_ctx, dst_buf, vk_tensor_offset(dst) + dst->view_offs, src_buf, vk_tensor_offset(src) + src->view_offs, ggml_nbytes(src)); + return true; + } + + return false; +} + +static void ggml_backend_vk_synchronize(ggml_backend_t backend) { + VK_LOG_DEBUG("ggml_backend_vk_synchronize()"); + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + if(ctx->transfer_ctx.expired()) { + return; + } + + vk_context transfer_ctx = ctx->transfer_ctx.lock(); + + ggml_vk_ctx_end(transfer_ctx); + + for (auto& cpy : transfer_ctx->in_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + + ggml_vk_submit(transfer_ctx, ctx->fence); + VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_backend_vk_synchronize waitForFences"); + ctx->device->device.resetFences({ ctx->fence }); + + for (auto& cpy : transfer_ctx->out_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + + ctx->transfer_ctx.reset(); +} + +static bool ggml_vk_is_empty(ggml_tensor * node) { + return ggml_is_empty(node) || node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE; +} + +static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { + VK_LOG_DEBUG("ggml_backend_vk_graph_compute(" << cgraph->n_nodes << " nodes)"); + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_vk_build_graph(ctx, cgraph->nodes[i], i, nullptr, 0, true, false, false); + } + ggml_vk_preallocate_buffers(ctx); + ggml_pipeline_allocate_descriptor_sets(ctx->device); + + int last_node = cgraph->n_nodes - 1; + + // If the last op in the cgraph isn't backend GPU, the command buffer doesn't get closed properly + while (last_node > 0 && ggml_vk_is_empty(cgraph->nodes[last_node])) { + last_node -= 1; + } + + // Reserve tensor context space for all nodes + ctx->tensor_ctxs.resize(cgraph->n_nodes); + + bool first_node_in_batch = true; // true if next node will be first node in a batch + int submit_node_idx = 0; // index to first node in a batch + + // Submit work every nodes_per_submit nodes to overlap CPU cmdbuffer generation with GPU execution. + // Start with a smaller count to get work submitted right away, and increase it after each submit. + int nodes_per_submit = 20; + int submitted_nodes = 0; + int submit_count = 0; + for (int i = 0; i < cgraph->n_nodes; i++) { + if (first_node_in_batch) { + submit_node_idx = i; + } + + bool submit = (submitted_nodes >= nodes_per_submit) || (i == last_node); + + bool enqueued = ggml_vk_build_graph(ctx, cgraph->nodes[i], i, cgraph->nodes[submit_node_idx], submit_node_idx, false, i == last_node, submit); + + if (enqueued) { + ++submitted_nodes; + +#ifndef GGML_VULKAN_CHECK_RESULTS + if (first_node_in_batch) { + first_node_in_batch = false; + } +#endif + } + + if (submit) { + first_node_in_batch = true; + submitted_nodes = 0; + switch (submit_count) { + case 0: + nodes_per_submit = 50; + break; + default: + nodes_per_submit = 100; + break; + } + submit_count++; + } + } + +#ifdef GGML_VULKAN_PERF + ctx->device->perf_logger->print_timings(); +#endif + + ggml_vk_graph_cleanup(ctx); + + return GGML_STATUS_SUCCESS; + + UNUSED(backend); +} + +// TODO: enable async and synchronize +static ggml_backend_i ggml_backend_vk_interface = { + /* .get_name = */ ggml_backend_vk_name, + /* .free = */ ggml_backend_vk_free, + /* .set_tensor_async = */ NULL, // ggml_backend_vk_set_tensor_async, + /* .get_tensor_async = */ NULL, // ggml_backend_vk_get_tensor_async, + /* .cpy_tensor_async = */ NULL, // ggml_backend_vk_cpy_tensor_async, + /* .synchronize = */ NULL, // ggml_backend_vk_synchronize, + /* .graph_plan_create = */ NULL, + /* .graph_plan_free = */ NULL, + /* .graph_plan_update = */ NULL, + /* .graph_plan_compute = */ NULL, + /* .graph_compute = */ ggml_backend_vk_graph_compute, + /* .event_record = */ NULL, + /* .event_wait = */ NULL, +}; + +static ggml_guid_t ggml_backend_vk_guid() { + static ggml_guid guid = { 0xb8, 0xf7, 0x4f, 0x86, 0x40, 0x3c, 0xe1, 0x02, 0x91, 0xc8, 0xdd, 0xe9, 0x02, 0x3f, 0xc0, 0x2b }; + return &guid; +} + +ggml_backend_t ggml_backend_vk_init(size_t dev_num) { + VK_LOG_DEBUG("ggml_backend_vk_init(" << dev_num << ")"); + + ggml_backend_vk_context * ctx = new ggml_backend_vk_context; + ggml_vk_init(ctx, dev_num); + + ggml_backend_t vk_backend = new ggml_backend { + /* .guid = */ ggml_backend_vk_guid(), + /* .interface = */ ggml_backend_vk_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), dev_num), + /* .context = */ ctx, + }; + + return vk_backend; +} + +bool ggml_backend_is_vk(ggml_backend_t backend) { + return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_vk_guid()); +} + +int ggml_backend_vk_get_device_count() { + return ggml_vk_get_device_count(); +} + +void ggml_backend_vk_get_device_description(int device, char * description, size_t description_size) { + GGML_ASSERT(device < (int) vk_instance.device_indices.size()); + int dev_idx = vk_instance.device_indices[device]; + ggml_vk_get_device_description(dev_idx, description, description_size); +} + +void ggml_backend_vk_get_device_memory(int device, size_t * free, size_t * total) { + GGML_ASSERT(device < (int) vk_instance.device_indices.size()); + + vk::PhysicalDevice vkdev = vk_instance.instance.enumeratePhysicalDevices()[vk_instance.device_indices[device]]; + + vk::PhysicalDeviceMemoryProperties memprops = vkdev.getMemoryProperties(); + + for (const vk::MemoryHeap& heap : memprops.memoryHeaps) { + if (heap.flags & vk::MemoryHeapFlagBits::eDeviceLocal) { + *total = heap.size; + *free = heap.size; + break; + } + } +} + +////////////////////////// + +struct ggml_backend_vk_device_context { + size_t device; + std::string name; + std::string description; +}; + +static const char * ggml_backend_vk_device_get_name(ggml_backend_dev_t dev) { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + return ctx->name.c_str(); +} + +static const char * ggml_backend_vk_device_get_description(ggml_backend_dev_t dev) { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + return ctx->description.c_str(); +} + +static void ggml_backend_vk_device_get_memory(ggml_backend_dev_t device, size_t * free, size_t * total) { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)device->context; + ggml_backend_vk_get_device_memory(ctx->device, free, total); +} + +static ggml_backend_buffer_type_t ggml_backend_vk_device_get_buffer_type(ggml_backend_dev_t dev) { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + return ggml_backend_vk_buffer_type(ctx->device); +} + +static ggml_backend_buffer_type_t ggml_backend_vk_device_get_host_buffer_type(ggml_backend_dev_t dev) { + UNUSED(dev); + return ggml_backend_vk_host_buffer_type(); +} + +static enum ggml_backend_dev_type ggml_backend_vk_device_get_type(ggml_backend_dev_t dev) { + UNUSED(dev); + return GGML_BACKEND_DEVICE_TYPE_GPU; +} + +static void ggml_backend_vk_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props) { + props->name = ggml_backend_vk_device_get_name(dev); + props->description = ggml_backend_vk_device_get_description(dev); + props->type = ggml_backend_vk_device_get_type(dev); + ggml_backend_vk_device_get_memory(dev, &props->memory_free, &props->memory_total); + props->caps = { + /* .async = */ false, + /* .host_buffer = */ true, + /* .buffer_from_host_ptr = */ false, + /* .events = */ false, + }; +} + +static ggml_backend_t ggml_backend_vk_device_init(ggml_backend_dev_t dev, const char * params) { + UNUSED(params); + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + return ggml_backend_vk_init(ctx->device); +} + +static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + switch (op->op) { + case GGML_OP_UNARY: + switch (ggml_get_unary_op(op)) { + case GGML_UNARY_OP_GELU: + case GGML_UNARY_OP_GELU_QUICK: + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_TANH: + return ggml_is_contiguous(op->src[0]); + default: + return false; + } + break; + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + const vk_device& device = ggml_vk_get_device(ctx->device); + if (op->op == GGML_OP_MUL_MAT_ID && !device->mul_mat_id_s && !device->mul_mat_id_m && !device->mul_mat_id_l) { + // If there's not enough shared memory for row_ids and the result tile, fallback to CPU + return false; + } + switch (op->src[0]->type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return false; + } + struct ggml_tensor * a; + struct ggml_tensor * b; + if (op->op == GGML_OP_MUL_MAT) { + a = op->src[0]; + b = op->src[1]; + } else { + a = op->src[2]; + b = op->src[1]; + } + if (a->ne[3] != b->ne[3]) { + return false; + } + if (!(ggml_vk_dim01_contiguous(op->src[0]) || op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) || + !(ggml_vk_dim01_contiguous(op->src[1]) || op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16)) { + return false; + } + + return true; + } break; + case GGML_OP_FLASH_ATTN_EXT: + { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + if (!ggml_vk_get_device(ctx->device)->coopmat2) { + return false; + } + switch (op->src[0]->ne[0]) { + case 64: + case 80: + case 96: + case 112: + case 128: + case 256: + break; + default: + return false; + } + if (op->src[0]->type != GGML_TYPE_F32) { + return false; + } + if (op->type != GGML_TYPE_F32) { + return false; + } + if (op->src[3] && op->src[3]->type != GGML_TYPE_F16) { + return false; + } + // It's straightforward to support different K/V dequant, but would + // significantly increase the number of pipelines + if (op->src[1]->type != op->src[2]->type) { + return false; + } + switch (op->src[1]->type) { + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + // K dequants currently disabled because D dimension is rounded up to 256 and runs inefficiently + //case GGML_TYPE_Q2_K: + //case GGML_TYPE_Q3_K: + //case GGML_TYPE_Q4_K: + //case GGML_TYPE_Q5_K: + //case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return false; + } + return true; + } + case GGML_OP_GET_ROWS: + { + switch (op->src[0]->type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_IQ4_NL: + return true; + default: + return false; + } + } break; + case GGML_OP_CONT: + case GGML_OP_CPY: + case GGML_OP_DUP: + { + ggml_type src0_type = op->src[0]->type; + ggml_type src1_type = op->src[1] != nullptr ? op->src[1]->type : src0_type; + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F16) { + return true; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return true; + } + return false; + } break; + case GGML_OP_REPEAT: + return ggml_type_size(op->type) == sizeof(float) && ggml_type_size(op->src[0]->type) == sizeof(float); + case GGML_OP_ROPE: + { + const int mode = ((const int32_t *) op->op_params)[2]; + if (mode & GGML_ROPE_TYPE_MROPE) { + return false; + } + if (mode & GGML_ROPE_TYPE_VISION) { + return false; + } + return ggml_is_contiguous(op->src[0]); + } + case GGML_OP_NONE: + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_NORM: + case GGML_OP_GROUP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_ADD: + case GGML_OP_ACC: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_SOFT_MAX: + case GGML_OP_ARGSORT: + case GGML_OP_SUM_ROWS: + case GGML_OP_IM2COL: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_POOL_2D: + case GGML_OP_RWKV_WKV6: + case GGML_OP_LEAKY_RELU: + return true; + default: + return false; + } + + UNUSED(dev); +} + +static bool ggml_backend_vk_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { + if (buft->iface.get_name != ggml_backend_vk_buffer_type_name) { + return false; + } + + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + ggml_backend_vk_buffer_type_context * buft_ctx = (ggml_backend_vk_buffer_type_context *)buft->context; + + return buft_ctx->device->idx == ctx->device; +} + +static bool ggml_backend_vk_device_offload_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + const int min_batch_size = 32; + + return (op->ne[1] >= min_batch_size && op->op != GGML_OP_GET_ROWS) || + (op->ne[2] >= min_batch_size && op->op == GGML_OP_MUL_MAT_ID); + + UNUSED(dev); +} + +static const struct ggml_backend_device_i ggml_backend_vk_device_i = { + /* .get_name = */ ggml_backend_vk_device_get_name, + /* .get_description = */ ggml_backend_vk_device_get_description, + /* .get_memory = */ ggml_backend_vk_device_get_memory, + /* .get_type = */ ggml_backend_vk_device_get_type, + /* .get_props = */ ggml_backend_vk_device_get_props, + /* .init_backend = */ ggml_backend_vk_device_init, + /* .get_buffer_type = */ ggml_backend_vk_device_get_buffer_type, + /* .get_host_buffer_type = */ ggml_backend_vk_device_get_host_buffer_type, + /* .buffer_from_host_ptr = */ NULL, + /* .supports_op = */ ggml_backend_vk_device_supports_op, + /* .supports_buft = */ ggml_backend_vk_device_supports_buft, + /* .offload_op = */ ggml_backend_vk_device_offload_op, + /* .event_new = */ NULL, + /* .event_free = */ NULL, + /* .event_synchronize = */ NULL, +}; + +static const char * ggml_backend_vk_reg_get_name(ggml_backend_reg_t reg) { + UNUSED(reg); + return GGML_VK_NAME; +} + +static size_t ggml_backend_vk_reg_get_device_count(ggml_backend_reg_t reg) { + UNUSED(reg); + return ggml_backend_vk_get_device_count(); +} + +static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg, size_t device) { + static std::vector devices; + + static bool initialized = false; + + { + static std::mutex mutex; + std::lock_guard lock(mutex); + if (!initialized) { + for (int i = 0; i < ggml_backend_vk_get_device_count(); i++) { + ggml_backend_vk_device_context * ctx = new ggml_backend_vk_device_context; + char desc[256]; + ggml_backend_vk_get_device_description(i, desc, sizeof(desc)); + ctx->device = i; + ctx->name = GGML_VK_NAME + std::to_string(i); + ctx->description = desc; + devices.push_back(new ggml_backend_device { + /* .iface = */ ggml_backend_vk_device_i, + /* .reg = */ reg, + /* .context = */ ctx, + }); + } + initialized = true; + } + } + + GGML_ASSERT(device < devices.size()); + return devices[device]; +} + +static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = { + /* .get_name = */ ggml_backend_vk_reg_get_name, + /* .get_device_count = */ ggml_backend_vk_reg_get_device_count, + /* .get_device = */ ggml_backend_vk_reg_get_device, + /* .get_proc_address = */ NULL, +}; + +ggml_backend_reg_t ggml_backend_vk_reg() { + static ggml_backend_reg reg = { + /* .api_version = */ GGML_BACKEND_API_VERSION, + /* .iface = */ ggml_backend_vk_reg_i, + /* .context = */ nullptr, + }; + + return ® +} + +// Extension availability +static bool ggml_vk_instance_validation_ext_available(const std::vector& instance_extensions) { +#ifdef GGML_VULKAN_VALIDATE + bool portability_enumeration_ext = false; + // Check for portability enumeration extension for MoltenVK support + for (const auto& properties : instance_extensions) { + if (strcmp("VK_KHR_portability_enumeration", properties.extensionName) == 0) { + return true; + } + } + if (!portability_enumeration_ext) { + std::cerr << "ggml_vulkan: WARNING: Instance extension VK_KHR_portability_enumeration not found." << std::endl; + } +#endif + return false; + + UNUSED(instance_extensions); +} +static bool ggml_vk_instance_portability_enumeration_ext_available(const std::vector& instance_extensions) { +#ifdef __APPLE__ + bool portability_enumeration_ext = false; + // Check for portability enumeration extension for MoltenVK support + for (const auto& properties : instance_extensions) { + if (strcmp("VK_KHR_portability_enumeration", properties.extensionName) == 0) { + return true; + } + } + if (!portability_enumeration_ext) { + std::cerr << "ggml_vulkan: WARNING: Instance extension VK_KHR_portability_enumeration not found." << std::endl; + } +#endif + return false; + + UNUSED(instance_extensions); +} + +static bool ggml_vk_khr_cooperative_matrix_support(const vk::PhysicalDeviceProperties& props, const vk::PhysicalDeviceDriverProperties& driver_props) { + switch (props.vendorID) { + case VK_VENDOR_ID_INTEL: + // Intel drivers don't support coopmat properly yet + return false; + case VK_VENDOR_ID_AMD: + if (driver_props.driverID == vk::DriverId::eAmdProprietary || driver_props.driverID == vk::DriverId::eAmdOpenSource) { + // Workaround for AMD proprietary driver reporting support on all GPUs + const std::string name = props.deviceName; + return name.rfind("AMD Radeon RX 7", 0) == 0 || name.rfind("AMD Radeon(TM) RX 7", 0) == 0 || // RDNA 3 consumer GPUs + name.rfind("AMD Radeon PRO W7", 0) == 0 || name.rfind("AMD Radeon(TM) PRO W7", 0) == 0 || // RDNA 3 workstation GPUs + name.rfind("AMD Radeon 7", 0) == 0 || name.rfind("AMD Radeon(TM) 7", 0) == 0; // RDNA 3 APUs + } + return true; + default: + return true; + } +} + +// checks + +#ifdef GGML_VULKAN_CHECK_RESULTS +static void ggml_vk_print_graph_origin(const ggml_tensor * tensor, std::vector& done, int level = 0) { + if (std::find(done.begin(), done.end(), tensor) != done.end() || level > 10) { + return; + } + for (int j = 0; j < level; j++) { + std::cerr << " "; + } + std::cerr << ggml_op_name(tensor->op) << " gpu=" << (tensor->extra != nullptr) << std::endl; + + done.push_back(tensor); + + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (tensor->src[i] != nullptr) { + ggml_vk_print_graph_origin(tensor->src[i], done, level + 1); + } + } +} + +static void ggml_vk_print_tensor_area(const ggml_tensor * tensor, const void * data, int i0, int i1, int i2, int i3) { + if (tensor->type != GGML_TYPE_F32 && tensor->type != GGML_TYPE_F16 && tensor->type != GGML_TYPE_I32) { + return; + } + i0 = std::max(i0, 5); + i1 = std::max(i1, 5); + i2 = std::max(i2, 0); + i3 = std::max(i3, 0); + fprintf(stderr, " "); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + fprintf(stderr, "%7d ", idx1); + } + fprintf(stderr, "\n"); + for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { + fprintf(stderr, "%7d: ", idx0); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + if (idx0 >= 0 && idx0 < tensor->ne[0] && idx1 >= 0 && idx1 < tensor->ne[1] && i2 >= 0 && i2 < tensor->ne[2] && i3 >= 0 && i3 < tensor->ne[3]) { + float val; + if (tensor->type == GGML_TYPE_F32) { + val = *(const float *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); + } else if (tensor->type == GGML_TYPE_F16) { + val = ggml_fp16_to_fp32(*(const ggml_fp16_t *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0])); + } else if (tensor->type == GGML_TYPE_I32) { + val = *(const int32_t *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); + } else { + GGML_ABORT("fatal error"); + } + fprintf(stderr, "% 7.2f ", val); + } else { + fprintf(stderr, " "); + } + } + fprintf(stderr, "\n"); + } +} + +static void ggml_vk_print_tensor(const ggml_tensor * tensor, const char * name) { + void * tensor_data = tensor->data; + + const bool is_gpu = tensor->buffer != nullptr && ggml_backend_buffer_is_vk(tensor->buffer); + + if (is_gpu) { + const size_t tensor_size = ggml_nbytes(tensor); + tensor_data = malloc(tensor_size); + + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; + + vk_buffer buffer_gpu = buf_ctx->dev_buffer; + ggml_vk_buffer_read(buffer_gpu, vk_tensor_offset(tensor) + tensor->view_offs, tensor_data, tensor_size); + } + + std::cerr << "TENSOR CHECK " << name << " (" << tensor->name << "): " << ggml_op_name(tensor->op) << std::endl; + std::cerr << "tensor=" << tensor << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << std::endl; + if (tensor->src[0] != nullptr) { + std::cerr << "tensor->src[0]=" << tensor->src[0] << " name=" << tensor->src[0]->name << " op=" << ggml_op_name(tensor->src[0]->op) << " type=" << ggml_type_name(tensor->src[0]->type) << " ne0=" << tensor->src[0]->ne[0] << " nb0=" << tensor->src[0]->nb[0] << " ne1=" << tensor->src[0]->ne[1] << " nb1=" << tensor->src[0]->nb[1] << " ne2=" << tensor->src[0]->ne[2] << " nb2=" << tensor->src[0]->nb[2] << " ne3=" << tensor->src[0]->ne[3] << " nb3=" << tensor->src[0]->nb[3] << std::endl; + } + if (tensor->src[1] != nullptr) { + std::cerr << "tensor->src[1]=" << tensor->src[1] << " name=" << tensor->src[1]->name << " op=" << ggml_op_name(tensor->src[1]->op) << " type=" << ggml_type_name(tensor->src[1]->type) << " ne0=" << tensor->src[1]->ne[0] << " nb0=" << tensor->src[1]->nb[0] << " ne1=" << tensor->src[1]->ne[1] << " nb1=" << tensor->src[1]->nb[1] << " ne2=" << tensor->src[1]->ne[2] << " nb2=" << tensor->src[1]->nb[2] << " ne3=" << tensor->src[1]->ne[3] << " nb3=" << tensor->src[1]->nb[3] << std::endl; + } + std::cerr << std::endl << "Result:" << std::endl; + ggml_vk_print_tensor_area(tensor, tensor_data, 5, 5, 0, 0); + std::cerr << std::endl; + std::vector done; + ggml_vk_print_graph_origin(tensor, done); + + if (is_gpu) { + free(tensor_data); + } +} + +void * comp_result; +size_t comp_size; +size_t comp_nb[GGML_MAX_DIMS]; +size_t check_counter = 0; +static void ggml_vk_check_results_0(ggml_tensor * tensor) { + if (tensor->op == GGML_OP_TRANSPOSE) { + return; + } + + check_counter++; + if (!(vk_output_tensor > 0 && vk_output_tensor == check_counter) && check_counter <= vk_skip_checks) { + return; + } + + VK_LOG_DEBUG("ggml_vk_check_results_0(" << tensor->name << ")"); + + ggml_tensor * src0 = tensor->src[0]; + ggml_tensor * src1 = tensor->src[1]; + ggml_tensor * src2 = tensor->src[2]; + ggml_tensor * src3 = tensor->src[3]; + + struct ggml_init_params iparams = { + /*.mem_size =*/ 2ul*1024ul*1024ul*1024ul, + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ false, + }; + + struct ggml_context * ggml_ctx = ggml_init(iparams); + + struct ggml_tensor * src0_clone = nullptr; + struct ggml_tensor * src1_clone = nullptr; + struct ggml_tensor * src2_clone = nullptr; + struct ggml_tensor * src3_clone = nullptr; + struct ggml_tensor * tensor_clone = nullptr; + + size_t src0_size; + size_t src1_size; + size_t src2_size; + size_t src3_size; + + void * src0_buffer = nullptr; + void * src1_buffer = nullptr; + void * src2_buffer = nullptr; + void * src3_buffer = nullptr; + + if (src0 != nullptr) { + src0_clone = ggml_dup_tensor(ggml_ctx, src0); + + src0_size = ggml_nbytes(src0); + + src0_buffer = malloc(src0_size); + src0_clone->data = src0_buffer; + if (ggml_backend_buffer_is_host(src0->buffer)) { + memcpy(src0_clone->data, src0->data, src0_size); + memcpy(src0_clone->nb, src0->nb, sizeof(size_t) * GGML_MAX_DIMS); + } else if (ggml_backend_buffer_is_vk(src0->buffer)) { + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + vk_buffer& buffer_gpu = buf_ctx->dev_buffer; + uint64_t offset = vk_tensor_offset(src0) + src0->view_offs; + if (!ggml_is_contiguous(src0) && ggml_vk_dim01_contiguous(src0)) { + for (int i3 = 0; i3 < src0->ne[3]; i3++) { + for (int i2 = 0; i2 < src0->ne[2]; i2++) { + const int idx = i3*src0->ne[2] + i2; + ggml_vk_buffer_read(buffer_gpu, offset + idx * src0->nb[2], ((char *)src0_clone->data + idx * src0_clone->nb[2]), src0->ne[1] * src0->nb[1]); + } + } + + src0_clone->nb[0] = src0->nb[0]; + src0_clone->nb[1] = src0->nb[1]; + for (int i = 2; i < GGML_MAX_DIMS; i++) { + src0_clone->nb[i] = src0_clone->nb[i - 1]*src0_clone->ne[i - 1]; + } + } else { + if (offset + src0_size >= buffer_gpu->size) { + src0_size = buffer_gpu->size - offset; + } + ggml_vk_buffer_read(buffer_gpu, offset, src0_clone->data, src0_size); + memcpy(src0_clone->nb, src0->nb, sizeof(size_t) * GGML_MAX_DIMS); + } + } else { + GGML_ABORT("fatal error"); + } + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + ggml_vk_print_tensor(src0, "src0"); + } + } + if (src1 != nullptr) { + src1_clone = ggml_dup_tensor(ggml_ctx, src1); + + src1_size = ggml_nbytes(src1); + + src1_buffer = malloc(src1_size); + src1_clone->data = src1_buffer; + if (ggml_backend_buffer_is_host(src1->buffer)) { + memcpy(src1_clone->data, src1->data, src1_size); + memcpy(src1_clone->nb, src1->nb, sizeof(size_t) * GGML_MAX_DIMS); + } else if (ggml_backend_buffer_is_vk(src1->buffer)) { + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + vk_buffer& buffer_gpu = buf_ctx->dev_buffer; + uint64_t offset = vk_tensor_offset(src1) + src1->view_offs; + if (!ggml_is_contiguous(src1) && ggml_vk_dim01_contiguous(src1)) { + for (int i3 = 0; i3 < src1->ne[3]; i3++) { + for (int i2 = 0; i2 < src1->ne[2]; i2++) { + const int idx = i3*src1->ne[2] + i2; + ggml_vk_buffer_read(buffer_gpu, offset + idx * src1->nb[2], ((char *)src1_clone->data + idx * src1_clone->nb[2]), src1->ne[1] * src1->nb[1]); + } + } + + src1_clone->nb[0] = src1->nb[0]; + src1_clone->nb[1] = src1->nb[1]; + for (int i = 2; i < GGML_MAX_DIMS; i++) { + src1_clone->nb[i] = src1_clone->nb[i - 1]*src1_clone->ne[i - 1]; + } + } else { + if (offset + src1_size >= buffer_gpu->size) { + src1_size = buffer_gpu->size - offset; + } + ggml_vk_buffer_read(buffer_gpu, offset, src1_clone->data, src1_size); + memcpy(src1_clone->nb, src1->nb, sizeof(size_t) * GGML_MAX_DIMS); + } + } else { + GGML_ABORT("fatal error"); + } + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + ggml_vk_print_tensor(src1, "src1"); + } + } + if (src2 != nullptr) { + src2_clone = ggml_dup_tensor(ggml_ctx, src2); + + src2_size = ggml_nbytes(src2); + + src2_buffer = malloc(src2_size); + src2_clone->data = src2_buffer; + if (ggml_backend_buffer_is_host(src2->buffer)) { + memcpy(src2_clone->data, src2->data, src2_size); + memcpy(src2_clone->nb, src2->nb, sizeof(size_t) * GGML_MAX_DIMS); + } else if (ggml_backend_buffer_is_vk(src2->buffer)) { + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src2->buffer->context; + vk_buffer& buffer_gpu = buf_ctx->dev_buffer; + uint64_t offset = vk_tensor_offset(src2) + src2->view_offs; + if (!ggml_is_contiguous(src2) && ggml_vk_dim01_contiguous(src2)) { + for (int i3 = 0; i3 < src2->ne[3]; i3++) { + for (int i2 = 0; i2 < src2->ne[2]; i2++) { + const int idx = i3*src2->ne[2] + i2; + ggml_vk_buffer_read(buffer_gpu, offset + idx * src2->nb[2], ((char *)src2_clone->data + idx * src2_clone->nb[2]), src2->ne[1] * src2->nb[1]); + } + } + + src2_clone->nb[0] = src2->nb[0]; + src2_clone->nb[1] = src2->nb[1]; + for (int i = 2; i < GGML_MAX_DIMS; i++) { + src2_clone->nb[i] = src2_clone->nb[i - 1]*src2_clone->ne[i - 1]; + } + } else { + if (offset + src2_size >= buffer_gpu->size) { + src2_size = buffer_gpu->size - offset; + } + ggml_vk_buffer_read(buffer_gpu, offset, src2_clone->data, src2_size); + memcpy(src2_clone->nb, src2->nb, sizeof(size_t) * GGML_MAX_DIMS); + } + } else { + GGML_ABORT("fatal error"); + } + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + ggml_vk_print_tensor(src2, "src2"); + } + } + if (src3 != nullptr) { + src3_clone = ggml_dup_tensor(ggml_ctx, src3); + + src3_size = ggml_nbytes(src3); + + src3_buffer = malloc(src3_size); + src3_clone->data = src3_buffer; + if (ggml_backend_buffer_is_host(src3->buffer)) { + memcpy(src3_clone->data, src3->data, src3_size); + memcpy(src3_clone->nb, src3->nb, sizeof(size_t) * GGML_MAX_DIMS); + } else if (ggml_backend_buffer_is_vk(src3->buffer)) { + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src3->buffer->context; + vk_buffer& buffer_gpu = buf_ctx->dev_buffer; + uint64_t offset = vk_tensor_offset(src3) + src3->view_offs; + if (!ggml_is_contiguous(src3) && ggml_vk_dim01_contiguous(src3)) { + for (int i3 = 0; i3 < src3->ne[3]; i3++) { + for (int i2 = 0; i2 < src3->ne[2]; i2++) { + const int idx = i3*src3->ne[2] + i2; + ggml_vk_buffer_read(buffer_gpu, offset + idx * src3->nb[2], ((char *)src3_clone->data + idx * src3_clone->nb[2]), src3->ne[1] * src3->nb[1]); + } + } + + src3_clone->nb[0] = src3->nb[0]; + src3_clone->nb[1] = src3->nb[1]; + for (int i = 2; i < GGML_MAX_DIMS; i++) { + src3_clone->nb[i] = src3_clone->nb[i - 1]*src3_clone->ne[i - 1]; + } + } else { + if (offset + src3_size >= buffer_gpu->size) { + src3_size = buffer_gpu->size - offset; + } + ggml_vk_buffer_read(buffer_gpu, offset, src3_clone->data, src3_size); + memcpy(src3_clone->nb, src3->nb, sizeof(size_t) * GGML_MAX_DIMS); + } + } else { + GGML_ABORT("fatal error"); + } + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + ggml_vk_print_tensor(src3, "src3"); + } + } + + if (tensor->op == GGML_OP_FLASH_ATTN_EXT) { + const float *params = (const float *)tensor->op_params; + tensor_clone = ggml_flash_attn_ext(ggml_ctx, src0_clone, src1_clone, src2_clone, src3_clone, params[0], params[1], params[2]); + } else if (tensor->op == GGML_OP_MUL_MAT) { + tensor_clone = ggml_mul_mat(ggml_ctx, src0_clone, src1_clone); + } else if (tensor->op == GGML_OP_MUL_MAT_ID) { + tensor_clone = ggml_mul_mat_id(ggml_ctx, src0_clone, src1_clone, src2_clone); + } else if (tensor->op == GGML_OP_MUL) { + tensor_clone = ggml_mul(ggml_ctx, src0_clone, src1_clone); + } else if (tensor->op == GGML_OP_DIV) { + tensor_clone = ggml_div(ggml_ctx, src0_clone, src1_clone); + } else if (tensor->op == GGML_OP_CONCAT) { + tensor_clone = ggml_concat(ggml_ctx, src0_clone, src1_clone, *(int *)tensor->op_params); + } else if (tensor->op == GGML_OP_UPSCALE) { + tensor_clone = ggml_upscale_ext(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); + } else if (tensor->op == GGML_OP_SCALE) { + tensor_clone = ggml_scale(ggml_ctx, src0_clone, ((float *)tensor->op_params)[0]); + } else if (tensor->op == GGML_OP_SQR) { + tensor_clone = ggml_sqr(ggml_ctx, src0_clone); + } else if (tensor->op == GGML_OP_SIN) { + tensor_clone = ggml_sin(ggml_ctx, src0_clone); + } else if (tensor->op == GGML_OP_COS) { + tensor_clone = ggml_cos(ggml_ctx, src0_clone); + } else if (tensor->op == GGML_OP_CLAMP) { + tensor_clone = ggml_clamp(ggml_ctx, src0_clone, ((float *)tensor->op_params)[0], ((float *)tensor->op_params)[1]); + } else if (tensor->op == GGML_OP_PAD) { + tensor_clone = ggml_pad(ggml_ctx, src0_clone, tensor->ne[0] - src0_clone->ne[0], tensor->ne[1] - src0_clone->ne[1], tensor->ne[2] - src0_clone->ne[2], tensor->ne[3] - src0_clone->ne[3]); + } else if (tensor->op == GGML_OP_REPEAT) { + tensor_clone = ggml_repeat(ggml_ctx, src0_clone, tensor); + } else if (tensor->op == GGML_OP_ADD) { + tensor_clone = ggml_add(ggml_ctx, src0_clone, src1_clone); + } else if (tensor->op == GGML_OP_ACC) { + tensor_clone = ggml_acc(ggml_ctx, src0_clone, src1_clone, tensor->op_params[0], tensor->op_params[1], tensor->op_params[2], tensor->op_params[3]); + } else if (tensor->op == GGML_OP_NORM) { + tensor_clone = ggml_norm(ggml_ctx, src0_clone, *(float *)tensor->op_params); + } else if (tensor->op == GGML_OP_GROUP_NORM) { + tensor_clone = ggml_group_norm(ggml_ctx, src0_clone, *(int *)tensor->op_params, ((float *)tensor->op_params)[1]); + } else if (tensor->op == GGML_OP_RMS_NORM) { + tensor_clone = ggml_rms_norm(ggml_ctx, src0_clone, *(float *)tensor->op_params); + } else if (tensor->op == GGML_OP_SOFT_MAX) { + if (src1 != nullptr) { + tensor_clone = ggml_soft_max_ext(ggml_ctx, src0_clone, src1_clone, ((float *)tensor->op_params)[0], ((float *)tensor->op_params)[1]); + } else { + tensor_clone = ggml_soft_max(ggml_ctx, src0_clone); + } + } else if (tensor->op == GGML_OP_DIAG_MASK_INF) { + tensor_clone = ggml_diag_mask_inf(ggml_ctx, src0_clone, *(int *)tensor->op_params); + } else if (tensor->op == GGML_OP_ROPE) { + const int n_dims = ((int32_t *) tensor->op_params)[1]; + const int mode = ((int32_t *) tensor->op_params)[2]; + //const int n_ctx_ggml = ((int32_t *) tensor->op_params)[3]; + const int n_ctx_orig_ggml = ((int32_t *) tensor->op_params)[4]; + const float freq_base = ((float *) tensor->op_params)[5]; + const float freq_scale = ((float *) tensor->op_params)[6]; + const float ext_factor = ((float *) tensor->op_params)[7]; + const float attn_factor = ((float *) tensor->op_params)[8]; + const float beta_fast = ((float *) tensor->op_params)[9]; + const float beta_slow = ((float *) tensor->op_params)[10]; + tensor_clone = ggml_rope_ext(ggml_ctx, src0_clone, src1_clone, src2_clone, n_dims, mode, n_ctx_orig_ggml, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + } else if (tensor->op == GGML_OP_UNARY) { + switch (ggml_get_unary_op(tensor)) { + case GGML_UNARY_OP_SILU: + tensor_clone = ggml_silu(ggml_ctx, src0_clone); + break; + case GGML_UNARY_OP_GELU: + tensor_clone = ggml_gelu(ggml_ctx, src0_clone); + break; + case GGML_UNARY_OP_GELU_QUICK: + tensor_clone = ggml_gelu_quick(ggml_ctx, src0_clone); + break; + case GGML_UNARY_OP_RELU: + tensor_clone = ggml_relu(ggml_ctx, src0_clone); + break; + case GGML_UNARY_OP_TANH: + tensor_clone = ggml_tanh(ggml_ctx, src0_clone); + break; + default: + std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl; + GGML_ABORT("fatal error"); + } + } else if (tensor->op == GGML_OP_CPY || tensor->op == GGML_OP_DUP) { + if (src1 == nullptr) { + tensor_clone = ggml_dup(ggml_ctx, src0_clone); + tensor_clone->type = tensor->type; + } else { + tensor_clone = ggml_cpy(ggml_ctx, src0_clone, src1_clone); + } + } else if (tensor->op == GGML_OP_CONT) { + tensor_clone = ggml_cont_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); + } else if (tensor->op == GGML_OP_RESHAPE) { + tensor_clone = ggml_reshape_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); + } else if (tensor->op == GGML_OP_VIEW) { + tensor_clone = ggml_view_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3], tensor->nb[1], tensor->nb[2], tensor->nb[3], ((int32_t *) tensor->op_params)[0]); + } else if (tensor->op == GGML_OP_PERMUTE) { + int32_t * params = (int32_t *)tensor->op_params; + tensor_clone = ggml_permute(ggml_ctx, src0_clone, params[0], params[1], params[2], params[3]); + } else if (tensor->op == GGML_OP_TRANSPOSE) { + tensor_clone = ggml_transpose(ggml_ctx, src0_clone); + } else if (tensor->op == GGML_OP_GET_ROWS) { + tensor_clone = ggml_get_rows(ggml_ctx, src0_clone, src1_clone); + } else if (tensor->op == GGML_OP_ARGSORT) { + tensor_clone = ggml_argsort(ggml_ctx, src0_clone, (ggml_sort_order) *(int *)tensor->op_params); + } else if (tensor->op == GGML_OP_SUM_ROWS) { + tensor_clone = ggml_sum_rows(ggml_ctx, src0_clone); + } else if (tensor->op == GGML_OP_IM2COL) { + const int32_t s0 = tensor->op_params[0]; + const int32_t s1 = tensor->op_params[1]; + const int32_t p0 = tensor->op_params[2]; + const int32_t p1 = tensor->op_params[3]; + const int32_t d0 = tensor->op_params[4]; + const int32_t d1 = tensor->op_params[5]; + + const bool is_2D = tensor->op_params[6] == 1; + tensor_clone = ggml_im2col(ggml_ctx, src0_clone, src1_clone, s0, s1, p0, p1, d0, d1, is_2D, tensor->type); + } else if (tensor->op == GGML_OP_TIMESTEP_EMBEDDING) { + const int32_t dim = tensor->op_params[0]; + const int32_t max_period = tensor->op_params[1]; + tensor_clone = ggml_timestep_embedding(ggml_ctx, src0_clone, dim, max_period); + } else if (tensor->op == GGML_OP_POOL_2D) { + enum ggml_op_pool op = static_cast(tensor->op_params[0]); + const int32_t k0 = tensor->op_params[1]; + const int32_t k1 = tensor->op_params[2]; + const int32_t s0 = tensor->op_params[3]; + const int32_t s1 = tensor->op_params[4]; + const int32_t p0 = tensor->op_params[5]; + const int32_t p1 = tensor->op_params[6]; + + tensor_clone = ggml_pool_2d(ggml_ctx, src0_clone, op, k0, k1, s0, s1, p0, p1); + } else if (tensor->op == GGML_OP_LEAKY_RELU) { + const float * op_params = (const float *)tensor->op_params; + tensor_clone = ggml_leaky_relu(ggml_ctx, src0_clone, op_params[0], false); + } else if (tensor->op == GGML_OP_RWKV_WKV6) { + tensor_clone = ggml_rwkv_wkv6(ggml_ctx, tensor->src[0], tensor->src[1], tensor->src[2], tensor->src[3], + tensor->src[4], tensor->src[5]); + } + else { + std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl; + GGML_ABORT("fatal error"); + } + + ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); + ggml_build_forward_expand(cgraph, tensor_clone); + + ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 8); + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + ggml_vk_print_tensor(tensor_clone, "tensor_clone"); + } + + comp_size = ggml_nbytes(tensor_clone); + + comp_result = malloc(comp_size); + memcpy(comp_result, tensor_clone->data, comp_size); + memcpy(comp_nb, tensor_clone->nb, sizeof(size_t) * GGML_MAX_DIMS); + + if (src0 != nullptr) { + free(src0_buffer); + } + if (src1 != nullptr) { + free(src1_buffer); + } + + ggml_free(ggml_ctx); + + VK_LOG_DEBUG("END ggml_vk_check_results_0(" << tensor->name << ")"); +} + +static void ggml_vk_check_results_1(ggml_tensor * tensor) { + if (tensor->op == GGML_OP_TRANSPOSE) { + return; + } + if (!(vk_output_tensor > 0 && vk_output_tensor == check_counter) && check_counter <= vk_skip_checks) { + return; + } + + VK_LOG_DEBUG("ggml_vk_check_results_1(" << tensor->name << ")"); + + ggml_tensor * src0 = tensor->src[0]; + ggml_tensor * src1 = tensor->src[1]; + ggml_tensor * src2 = tensor->src[2]; + + void * tensor_data = tensor->data; + + if (ggml_backend_buffer_is_vk(tensor->buffer)) { + size_t tensor_size = ggml_nbytes(tensor); + tensor_data = malloc(tensor_size); + + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; + + vk_buffer& buffer_gpu = buf_ctx->dev_buffer; + uint64_t offset = vk_tensor_offset(tensor) + tensor->view_offs; + if (offset + tensor_size >= buffer_gpu->size) { + tensor_size = buffer_gpu->size - offset; + } + + ggml_vk_buffer_read(buffer_gpu, offset, tensor_data, tensor_size); + } + + float first_error_result = -1.0f; + float first_error_correct = -1.0f; + std::array first_error = { -1, -1, -1, -1 }; + double avg_err = 0.0; + size_t counter = 0; + + for (int i3 = 0; i3 < tensor->ne[3]; i3++) { + for (int i2 = 0; i2 < tensor->ne[2]; i2++) { + for (int i1 = 0; i1 < tensor->ne[1]; i1++) { + for (int i0 = 0; i0 < tensor->ne[0]; i0++) { + const bool buffer_size_fit = i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0] < comp_size; + float correct = 0.0f; + float result = 0.0f; + + if (buffer_size_fit) { + if (tensor->type == GGML_TYPE_F32) { + correct = *(float *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0]); + result = *(float *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0]); + } else if (tensor->type == GGML_TYPE_F16) { + correct = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0])); + result = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0])); + } else if (tensor->type == GGML_TYPE_I32) { + correct = *(int32_t *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0]); + result = *(int32_t *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0]); + } else { + std::cerr << "Results check not implemented for type " << ggml_type_name(tensor->type) << std::endl; + } + } else { + std::cerr << "Missing debug code for type " << ggml_type_name(tensor->type) << std::endl; + GGML_ABORT("fatal error"); + } + + if ((std::isnan(correct) != std::isnan(result)) || (std::isinf(correct) != std::isinf(result)) || !buffer_size_fit) { + std::cerr << "ERROR: Invalid value in " << ggml_op_name(tensor->op) << " i3=" << i3 << " i2=" << i2 << " i1=" << i1 << " i0=" << i0 << " result=" << result << " correct=" << correct << " avg_err=" << (avg_err / counter) << std::endl; + std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; + if (src0 != nullptr) { + std::cerr << "src0=" << src0 << " src0->name=" << src0->name << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; + } + if (src1 != nullptr) { + std::cerr << "src1=" << src1 << " src1->name=" << src1->name << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; + } + if (src2 != nullptr) { + std::cerr << "src2=" << src2 << " src2->name=" << src2->name << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; + } + std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; + std::cerr << std::endl << "Result:" << std::endl; + ggml_vk_print_tensor_area(tensor, tensor_data, i0, i1, i2, i3); + std::cerr << std::endl << "Correct:" << std::endl; + ggml_vk_print_tensor_area(tensor, comp_result, i0, i1, i2, i3); + std::cerr << std::endl; + std::vector done; + ggml_vk_print_graph_origin(tensor, done); + GGML_ABORT("fatal error"); + } + if (first_error[0] == -1 && std::fabs(correct - result) > 0.1f) { + first_error[0] = i0; + first_error[1] = i1; + first_error[2] = i2; + first_error[3] = i3; + first_error_result = result; + first_error_correct = correct; + } + + // Special case, value is infinite, avoid NaN result in avg_err + // NaN also appears in results, if both are nan error is 0 + if (!std::isinf(correct) && !std::isinf(result) && !std::isnan(correct) && !std::isnan(result)) { + avg_err += std::fabs(correct - result); + } + counter++; + } + } + } + } + + avg_err /= counter; + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + std::cerr << "TENSOR CHECK: avg_err=" << avg_err << " in " << ggml_op_name(tensor->op) << " (check " << check_counter << ")" << std::endl; + std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; + if (src0 != nullptr) { + std::cerr << "src0=" << src0 << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; + } + if (src1 != nullptr) { + std::cerr << "src1=" << src1 << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; + } + if (src2 != nullptr) { + std::cerr << "src2=" << src2 << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; + } + std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; + std::cerr << std::endl << "Result:" << std::endl; + ggml_vk_print_tensor_area(tensor, tensor_data, 5, 5, 0, 0); + std::cerr << std::endl << "Correct:" << std::endl; + ggml_vk_print_tensor_area(tensor, comp_result, 5, 5, 0, 0); + std::cerr << std::endl; + std::vector done; + ggml_vk_print_graph_origin(tensor, done); + } + + if (avg_err > 0.05 || std::isnan(avg_err)) { + std::cerr << "ERROR: avg_err=" << avg_err << " in " << ggml_op_name(tensor->op) << " (check " << check_counter << ")" << std::endl; + std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; + if (src0 != nullptr) { + std::cerr << "src0=" << src0 << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; + } + if (src1 != nullptr) { + std::cerr << "src1=" << src1 << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; + } + if (src2 != nullptr) { + std::cerr << "src2=" << src2 << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; + } + std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; + std::cerr << std::endl << "Result:" << std::endl; + ggml_vk_print_tensor_area(tensor, tensor_data, first_error[0], first_error[1], first_error[2], first_error[3]); + std::cerr << std::endl << "Correct:" << std::endl; + ggml_vk_print_tensor_area(tensor, comp_result, first_error[0], first_error[1], first_error[2], first_error[3]); + std::cerr << std::endl; + std::vector done; + ggml_vk_print_graph_origin(tensor, done); + GGML_ABORT("fatal error"); + } else { + std::cerr << check_counter << " " << tensor->name << " op=" << ggml_op_name(tensor->op) << " avg_err=" << avg_err << std::endl; + } + + free(comp_result); + comp_result = nullptr; + comp_size = 0; + + if (ggml_backend_buffer_is_vk(tensor->buffer)) { + free(tensor_data); + } + + VK_LOG_DEBUG("END ggml_vk_check_results_1(" << tensor->name << ")"); +} +#endif + +GGML_BACKEND_DL_IMPL(ggml_backend_vk_reg) diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt new file mode 100644 index 000000000..bd0c74cb1 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt @@ -0,0 +1,9 @@ +find_package (Threads REQUIRED) +find_package(Vulkan COMPONENTS glslc REQUIRED) + +set(TARGET vulkan-shaders-gen) +add_executable(${TARGET} vulkan-shaders-gen.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_compile_features(${TARGET} PRIVATE cxx_std_17) +target_link_libraries(vulkan-shaders-gen PUBLIC Threads::Threads) +target_link_libraries(vulkan-shaders-gen PRIVATE Vulkan::Vulkan) diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp new file mode 100644 index 000000000..d896f1ef0 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp @@ -0,0 +1,29 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = gl_GlobalInvocationID.x; + if (idx >= p.ne) { + return; + } + + const uint offset = p.param3; + const uint src1_i = idx - offset; + const uint oz = src1_i / p.nb02; + const uint oy = (src1_i - (oz * p.nb02)) / p.nb01; + const uint ox = src1_i % p.nb01; + + uint i00, i01, i02, i03; + get_indices(idx, i00, i01, i02, i03); + + if (ox < p.ne10 && oy < p.ne11 && oz < p.ne12) { + data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) + FLOAT_TYPE(data_b[get_boffset() + ox + oy * p.ne10 + oz * p.ne10 * p.ne11])); + } else { + data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)])); + } +} + diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp new file mode 100644 index 000000000..2b4085c4f --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp @@ -0,0 +1,29 @@ +#version 450 + +#extension GL_EXT_shader_16bit_storage : require + +#include "types.comp" +#include "generic_binary_head.comp" + +const uint num_threads = 256; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + uint idx = get_idx(); + + // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation + const uint num_iter = 2; + + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (idx >= p.ne) { + continue; + } + uint i00, i01, i02, i03; + get_indices(idx, i00, i01, i02, i03); + + data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) + FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); + + idx += num_threads; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp new file mode 100644 index 000000000..d4fa45b1e --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp @@ -0,0 +1,69 @@ +#version 450 + +#include "types.comp" + +#define BLOCK_SIZE 1024 +#define ASC 0 + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) buffer D {int data_d[];}; + +layout (push_constant) uniform parameter { + uint ncols; + uint ncols_pad; + uint order; +} p; + +shared int dst_row[BLOCK_SIZE]; + +void swap(uint idx0, uint idx1) { + int tmp = dst_row[idx0]; + dst_row[idx0] = dst_row[idx1]; + dst_row[idx1] = tmp; +} + +void main() { + // bitonic sort + const int col = int(gl_LocalInvocationID.x); + const uint row = gl_WorkGroupID.y; + + const uint row_offset = row * p.ncols; + + // initialize indices + if (col < p.ncols_pad) { + dst_row[col] = col; + } + barrier(); + + for (uint k = 2; k <= p.ncols_pad; k *= 2) { + for (uint j = k / 2; j > 0; j /= 2) { + const uint ixj = col ^ j; + if (col < p.ncols_pad && ixj > col) { + if ((col & k) == 0) { + if (dst_row[col] >= p.ncols || + (dst_row[ixj] < p.ncols && (p.order == ASC ? + data_a[row_offset + dst_row[col]] > data_a[row_offset + dst_row[ixj]] : + data_a[row_offset + dst_row[col]] < data_a[row_offset + dst_row[ixj]])) + ) { + swap(col, ixj); + } + } else { + if (dst_row[ixj] >= p.ncols || + (dst_row[col] < p.ncols && (p.order == ASC ? + data_a[row_offset + dst_row[col]] < data_a[row_offset + dst_row[ixj]] : + data_a[row_offset + dst_row[col]] > data_a[row_offset + dst_row[ixj]])) + ) { + swap(col, ixj); + } + } + } + barrier(); + } + } + + if (col < p.ncols) { + data_d[row_offset + col] = dst_row[col]; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp new file mode 100644 index 000000000..1e5cb8dae --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp @@ -0,0 +1,17 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + + const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(val < p.param1 ? p.param1 : (val > p.param2 ? p.param2 : val)); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp new file mode 100644 index 000000000..9ee2f1fae --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp @@ -0,0 +1,41 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + const int dim = p.param3; + + if (idx >= p.ne) { + return; + } + + const uint i3 = idx / (p.ne22*p.ne21*p.ne20); + const uint i3_offset = i3 * p.ne22*p.ne21*p.ne20; + const uint i2 = (idx - i3_offset) / (p.ne21*p.ne20); + const uint i2_offset = i2*p.ne21*p.ne20; + const uint i1 = (idx - i3_offset - i2_offset) / p.ne20; + const uint i0 = idx - i3_offset - i2_offset - i1*p.ne20; + + uint o[4] = {0, 0, 0, 0}; + o[dim] = dim == 0 ? p.ne00 : (dim == 1 ? p.ne01 : (dim == 2 ? p.ne02 : p.ne03)); + + const uint src0_idx = i3*p.nb03 + i2*p.nb02 + i1*p.nb01 + i0*p.nb00; + const uint src1_idx = (i3 - o[3])*p.nb13 + (i2 - o[2])*p.nb12 + (i1 - o[1])*p.nb11 + (i0 - o[0])*p.nb10; + const uint dst_idx = i3*p.nb23 + i2*p.nb22 + i1*p.nb21 + i0*p.nb20; + + const bool is_src0 = i0 < p.ne00 && i1 < p.ne01 && i2 < p.ne02 && i3 < p.ne03; + +#ifndef OPTIMIZATION_ERROR_WORKAROUND + data_d[get_doffset() + dst_idx] = D_TYPE(is_src0 ? data_a[get_aoffset() + src0_idx] : data_b[get_boffset() + src1_idx]); +#else + if (is_src0) { + data_d[get_doffset() + dst_idx] = data_a[get_aoffset() + src0_idx]; + } else { + data_d[get_doffset() + dst_idx] = data_b[get_boffset() + src1_idx]; + } +#endif +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp new file mode 100644 index 000000000..dd828c232 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp @@ -0,0 +1,42 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +#extension GL_EXT_control_flow_attributes : require + +const uint num_threads = 128; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + uint idx = get_idx(); + + // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation + const uint num_iter = 4; + + // fast path for when all four iterations are in-bounds + if (idx + (num_iter-1)*num_threads < p.ne) { + [[unroll]] for (uint i = 0; i < num_iter; ++i) { +#ifndef OPTIMIZATION_ERROR_WORKAROUND + data_d[get_doffset() + idx] = D_TYPE(data_a[get_aoffset() + idx]); +#else + data_d[get_doffset() + idx] = data_a[get_aoffset() + idx]; +#endif + idx += num_threads; + } + } else { + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (idx >= p.ne) { + continue; + } + +#ifndef OPTIMIZATION_ERROR_WORKAROUND + data_d[get_doffset() + idx] = D_TYPE(data_a[get_aoffset() + idx]); +#else + data_d[get_doffset() + idx] = data_a[get_aoffset() + idx]; +#endif + idx += num_threads; + } + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp new file mode 100644 index 000000000..29c906494 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp @@ -0,0 +1,20 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + +#ifndef OPTIMIZATION_ERROR_WORKAROUND + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(data_a[get_aoffset() + src0_idx(idx)]); +#else + data_d[get_doffset() + dst_idx(idx)] = data_a[get_aoffset() + src0_idx(idx)]; +#endif +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp new file mode 100644 index 000000000..0b8d02f58 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp @@ -0,0 +1,17 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + + const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(cos(val)); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp new file mode 100644 index 000000000..a4d3fca55 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp @@ -0,0 +1,20 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {float data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_GlobalInvocationID.x * 16; + + if (i >= p.nel) { + return; + } + + [[unroll]] for (uint l = 0; l < 16; l++) { + data_b[i + l] = D_TYPE(data_a[i + l]); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp new file mode 100644 index 000000000..91bb8f8db --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp @@ -0,0 +1,118 @@ +#if !defined(DATA_A_F32) && !defined(DATA_A_F16) +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#endif + +#include "types.comp" + +#if defined(A_TYPE_PACKED16) +layout (binding = 0) readonly buffer A_PACKED16 {A_TYPE_PACKED16 data_a_packed16[];}; +#endif +#if defined(A_TYPE_PACKED32) +layout (binding = 0) readonly buffer A_PACKED32 {A_TYPE_PACKED32 data_a_packed32[];}; +#endif + +#if defined(DATA_A_F32) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]); +} +#endif + +#if defined(DATA_A_F16) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]); +} +#endif + +#if defined(DATA_A_Q4_0) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + return (vec2(vui & 0xF, vui >> 4) - 8.0f); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); + return (vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, vui >> 12) - 8.0f); +} +#endif + +#if defined(DATA_A_Q4_1) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + return vec2(vui & 0xF, vui >> 4); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); + return vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, vui >> 12); +} +#endif + +#if defined(DATA_A_Q5_0) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint uint_qh = uint(data_a[a_offset + ib].qh[1]) << 16 | data_a[a_offset + ib].qh[0]; + const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + return (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint uint_qh = uint(data_a_packed16[a_offset + ib].qh[1]) << 16 | data_a_packed16[a_offset + ib].qh[0]; + const ivec2 qh0 = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const ivec2 qh1 = ivec2(((uint_qh >> (iqs + 1)) << 4) & 0x10, (uint_qh >> (iqs + 13)) & 0x10); + const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); + return (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f); +} +#endif + +#if defined(DATA_A_Q5_1) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint uint_qh = data_a[a_offset + ib].qh; + const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + return vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint uint_qh = data_a_packed16[a_offset + ib].qh; + const ivec2 qh0 = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const ivec2 qh1 = ivec2(((uint_qh >> (iqs + 1)) << 4) & 0x10, (uint_qh >> (iqs + 13)) & 0x10); + const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); + return vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y); +} +#endif + +#if defined(DATA_A_Q8_0) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + return vec2(int(data_a[a_offset + ib].qs[iqs]), int(data_a[a_offset + ib].qs[iqs + 1])); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + uint32_t v0 = data_a_packed16[a_offset + ib].qs[iqs/2]; + uint32_t v1 = data_a_packed16[a_offset + ib].qs[iqs/2 + 1]; + return vec4(int8_t(v0 & 0xFF), int8_t(v0 >> 8), int8_t(v1 & 0xFF), int8_t(v1 >> 8)); +} +#endif + +#if defined(DATA_A_IQ4_NL) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + return vec2(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[vui >> 4]); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); + return vec4(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[(vui >> 4) & 0xF], kvalues_iq4nl[(vui >> 8) & 0xF], kvalues_iq4nl[vui >> 12]); +} +#endif + +#if defined(DATA_A_F32) || defined(DATA_A_F16) +vec2 get_dm(uint ib, uint a_offset) { + return vec2(0, 0); +} +#endif + +#if defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ4_NL) +vec2 get_dm(uint ib, uint a_offset) { + return vec2(float(data_a[a_offset + ib].d), 0); +} +#endif + +#if defined(DATA_A_Q4_1) || defined(DATA_A_Q5_1) +vec2 get_dm(uint ib, uint a_offset) { + return vec2(float(data_a[a_offset + ib].d), float(data_a[a_offset + ib].m)); +} +#endif diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp new file mode 100644 index 000000000..94b78598e --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp @@ -0,0 +1,325 @@ + +#include "types.comp" + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ4_0 { + block_q4_0_packed16 block; +}; + +float16_t dequantFuncQ4_0(const in decodeBufQ4_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const uint idx = coordInBlock[1]; + const uint shift = (idx & 0x10) >> 2; + uint32_t qs = uint32_t(bl.block.qs[(idx & 0xE) >> 1]); + qs >>= shift; + qs &= 0x0F0F; + qs = unpack8(qs)[idx & 1]; + float16_t ret = (float16_t(qs) - float16_t(8)) * d; + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 4) buffer decodeBufQ4_1 { + block_q4_1 block; +}; + +float16_t dequantFuncQ4_1(const in decodeBufQ4_1 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const float16_t m = bl.block.m; + const uint idx = coordInBlock[1]; + const uint iqs = idx & 0xF; + const uint shift = (idx & 0x10) >> 2; + uint32_t qs = bl.block.qs[iqs]; + qs >>= shift; + qs &= 0xF; + float16_t ret = float16_t(qs) * d + m; + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ5_0 { + block_q5_0 block; +}; + +float16_t dequantFuncQ5_0(const in decodeBufQ5_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const uint idx = coordInBlock[1]; + const uint iqs = idx & 0xF; + + const uint uint_qh = uint(bl.block.qh[1]) << 16 | bl.block.qh[0]; + const uint qh = ((uint_qh >> idx) << 4) & 0x10; + + const uint shift = (idx & 0x10) >> 2; + uint32_t qs = bl.block.qs[iqs]; + qs >>= shift; + qs &= 0xF; + + float16_t ret = (float16_t(qs | qh) - float16_t(16)) * d; + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 8) buffer decodeBufQ5_1 { + block_q5_1 block; +}; + +float16_t dequantFuncQ5_1(const in decodeBufQ5_1 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const float16_t m = bl.block.m; + const uint idx = coordInBlock[1]; + const uint iqs = idx & 0xF; + + const uint uint_qh = bl.block.qh; + const uint qh = ((uint_qh >> idx) << 4) & 0x10; + + const uint shift = (idx & 0x10) >> 2; + uint32_t qs = bl.block.qs[iqs]; + qs >>= shift; + qs &= 0xF; + + float16_t ret = float16_t(qs | qh) * d + m; + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ8_0 { + block_q8_0_packed16 block; +}; + +float16_t dequantFuncQ8_0(const in decodeBufQ8_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const uint idx = coordInBlock[1]; + const uint iqs = idx; + + // Load 16b and select the byte for this element + int32_t qs = unpack8(int32_t(bl.block.qs[(iqs & 0x1E) >> 1]))[iqs & 1]; + float16_t ret = float16_t(qs) * d; + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 4) buffer decodeBufQ2_K { + block_q2_K block; +}; + +float16_t dequantFuncQ2_K(const in decodeBufQ2_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const f16vec2 d = bl.block.d; + const uint idx = coordInBlock[1]; + const uint iqs = idx; + + const uint qsi = (iqs / 128) * 32 + (iqs % 32); // 0..31 + const uint scalesi = iqs / 16; // 0..15 + const uint qsshift = ((iqs % 128) / 32) * 2; // 0,2,4,6 + + uint32_t qs = bl.block.qs[qsi]; + const uint scales = bl.block.scales[scalesi]; + float16_t ret = d.x * float16_t(scales & 0xF) * float16_t((qs >> qsshift) & 3) - d.y * float16_t(scales >> 4); + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ3_K { + block_q3_K block; +}; + +float16_t dequantFuncQ3_K(const in decodeBufQ3_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint idx = coordInBlock[1]; + const uint iqs = idx; + + const uint n = iqs / 128; // 0,1 + const uint qsi = n * 32 + (iqs % 32); // 0..63 + const uint hmi = (iqs % 32); // 0..31 + const uint j = (iqs % 128) / 8; // 0..15 + const uint is = iqs / 16; // 0..15 + const uint halfsplit = ((iqs % 128) / 32); // 0,1,2,3 + const uint qsshift = halfsplit * 2; // 0,2,4,6 + const uint m = 1 << (4 * n + halfsplit); // 1,2,4,8,16,32,64,128 + + uint32_t scaleidx0 = (is < 8) ? is : (is-8); + uint32_t scaleidx0shift = (is < 8) ? 0 : 4; + uint32_t scaleidx1 = is + 8 - (is/4)*4; + uint32_t scaleidx1shift = (is/4)*2; + + const int8_t us = int8_t(((bl.block.scales[scaleidx0] >> scaleidx0shift) & 0xF) | (((bl.block.scales[scaleidx1] >> scaleidx1shift) & 3) << 4)); + + const float16_t dl = bl.block.d * float16_t(us - 32); + + float16_t ret = dl * float16_t(int8_t((bl.block.qs[qsi ] >> qsshift) & 3) - (((bl.block.hmask[hmi ] & m) != 0) ? 0 : 4)); + + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ4_K { + block_q4_K block; +}; + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ4_K_packed16 { + block_q4_K_packed16 block; +}; + +float16_t dequantFuncQ4_K(const in decodeBufQ4_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + decodeBufQ4_K_packed16 bl16 = decodeBufQ4_K_packed16(bl); + const uint idx = coordInBlock[1]; + + const uint b = (idx & 0x20) >> 5; // 0,1 + const uint is = (idx & 0xE0) >> 5; // 0..7 + + const f16vec2 loadd = bl.block.d; + + uint32_t sc; + uint32_t mbyte; + + uint32_t scidx0 = (is < 4) ? is : (is + 4); + uint32_t scidx1 = (is < 4) ? is : (is - 4); + uint32_t scidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint32_t scidxshift1 = (is < 4) ? 0 : 2; + uint32_t mbidx0 = is + 4; + uint32_t mbidx1 = (is < 4) ? is + 4 : is; + uint32_t mbidxmask0 = (is < 4) ? 0xF : 0xF0; + uint32_t mbidxshift0 = (is < 4) ? 0 : 4; + uint32_t mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint32_t mbidxshift1 = (is < 4) ? 0 : 2; + + sc = uint8_t((bl.block.scales[scidx0] & 0xF) | ((bl.block.scales[scidx1] & scidxmask1) >> scidxshift1)); + mbyte = uint8_t(((bl.block.scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((bl.block.scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const float16_t d = loadd.x * float16_t(sc); + const float16_t m = loadd.y * float16_t(mbyte); + + uint qs = uint32_t(bl16.block.qs[((idx & 0xC0) >> 2) + ((idx & 0x1E) >> 1)]); + qs = (qs >> (b * 4)) & 0x0F0F; + qs = unpack8(qs)[idx & 1]; + + float16_t ret = d * float16_t(qs) - m; + + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ5_K { + block_q5_K block; +}; + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ5_K_packed16 { + block_q5_K_packed16 block; +}; + +float16_t dequantFuncQ5_K(const in decodeBufQ5_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + decodeBufQ5_K_packed16 bl16 = decodeBufQ5_K_packed16(bl); + const uint idx = coordInBlock[1]; + + const uint b = (idx & 0x20) >> 5; // 0,1 + const uint is = (idx & 0xE0) >> 5; // 0..7 + + const uint32_t hm = 0x0101 << is; + + const f16vec2 loadd = bl.block.d; + + uint32_t sc; + uint32_t mbyte; + + uint32_t scidx0 = (is < 4) ? is : (is + 4); + uint32_t scidx1 = (is < 4) ? is : (is - 4); + uint32_t scidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint32_t scidxshift1 = (is < 4) ? 0 : 2; + uint32_t mbidx0 = is + 4; + uint32_t mbidx1 = (is < 4) ? is + 4 : is; + uint32_t mbidxmask0 = (is < 4) ? 0xF : 0xF0; + uint32_t mbidxshift0 = (is < 4) ? 0 : 4; + uint32_t mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint32_t mbidxshift1 = (is < 4) ? 0 : 2; + + sc = uint8_t((bl.block.scales[scidx0] & 0xF) | ((bl.block.scales[scidx1] & scidxmask1) >> scidxshift1)); + mbyte = uint8_t(((bl.block.scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((bl.block.scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const float16_t d = loadd.x * float16_t(sc); + const float16_t m = loadd.y * float16_t(mbyte); + + uint qh = uint32_t(bl16.block.qh[(idx & 0x1E) >> 1]); + qh = qh & hm; + qh = unpack8(qh)[idx & 1]; + + uint qs = uint32_t(bl16.block.qs[((idx & 0xC0) >> 2) + ((idx & 0x1E) >> 1)]); + qs = (qs >> (b * 4)) & 0x0F0F; + qs = unpack8(qs)[idx & 1]; + + float16_t ret = d * (float16_t(qs) + (qh != 0 ? float16_t(16) : float16_t(0))) - m; + + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ6_K { + block_q6_K block; +}; + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ6_K_packed16 { + block_q6_K_packed16 block; +}; + +float16_t dequantFuncQ6_K(const in decodeBufQ6_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + decodeBufQ6_K_packed16 bl16 = decodeBufQ6_K_packed16(bl); + const uint idx = coordInBlock[1]; + + const uint b = (idx & 0x40) >> 6; // 0,1 + const uint qhshift = (idx & 0x60) >> 4; // 0,2,4,6 + const uint is = (idx & 0xF0) >> 4; // 0..15 + + const float16_t dscale = bl.block.d * float16_t(bl.block.scales[is]); + + uint ql = uint32_t(bl16.block.ql[((idx & 0x80) >> 2) + ((idx & 0x3E) >> 1)]); + ql = (ql >> (b * 4)) & 0x0F0F; + + uint qh = uint32_t(bl16.block.qh[((idx & 0x80) >> 3) + ((idx & 0x1E) >> 1)]); + qh = ((qh >> qhshift) & 0x0303) << 4; + + int q = unpack8(ql | qh)[idx & 1]; + + float16_t ret = dscale * float16_t(q - 32); + + return ret; +} + +#if defined(DATA_A_IQ4_NL) +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufIQ4_NL { + block_iq4_nl block; +}; + +float16_t dequantFuncIQ4_NL(const in decodeBufIQ4_NL bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const uint idx = coordInBlock[1]; + const uint iqs = idx & 0xF; + const uint shift = (idx & 0x10) >> 2; + uint32_t qs = bl.block.qs[iqs]; + qs >>= shift; + qs &= 0xF; + float16_t ret = float16_t(kvalues_iq4nl[qs]) * d; + return ret; +} +#endif + +#if defined(DATA_A_Q4_0) +#define dequantFuncA dequantFuncQ4_0 +#elif defined(DATA_A_Q4_1) +#define dequantFuncA dequantFuncQ4_1 +#elif defined(DATA_A_Q5_0) +#define dequantFuncA dequantFuncQ5_0 +#elif defined(DATA_A_Q5_1) +#define dequantFuncA dequantFuncQ5_1 +#elif defined(DATA_A_Q8_0) +#define dequantFuncA dequantFuncQ8_0 +#elif defined(DATA_A_Q2_K) +#define dequantFuncA dequantFuncQ2_K +#elif defined(DATA_A_Q3_K) +#define dequantFuncA dequantFuncQ3_K +#elif defined(DATA_A_Q4_K) +#define dequantFuncA dequantFuncQ4_K +#elif defined(DATA_A_Q5_K) +#define dequantFuncA dequantFuncQ5_K +#elif defined(DATA_A_Q6_K) +#define dequantFuncA dequantFuncQ6_K +#elif defined(DATA_A_IQ4_NL) +#define dequantFuncA dequantFuncIQ4_NL +#endif diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp new file mode 100644 index 000000000..8d806435b --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp @@ -0,0 +1,13 @@ +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require + +layout (push_constant) uniform parameter +{ + uint M; + uint K; + uint stride_a; + uint stride_b; + uint nel; +} p; + +#include "types.comp" diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp new file mode 100644 index 000000000..8de14fc03 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp @@ -0,0 +1,32 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_iq4_nl data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + init_iq4nl_shmem(); + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint q_idx = 8*il; + const uint b_idx = 1024*i + 32*ir + q_idx; + + const float d = float(data_a[ib].d); + + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_b[b_idx + l + 0] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] & 0xF]); + data_b[b_idx + l + 16] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] >> 4]); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp new file mode 100644 index 000000000..157154af3 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp @@ -0,0 +1,34 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { + const uint i = gl_WorkGroupID.x * 256 + wgy; + if (i >= p.M * p.K / QUANT_K) { + return; + } + + const uint tid = gl_LocalInvocationID.x; + const uint ip = tid / 32; + const uint il = tid - 32 * ip; + const uint is = 8 * ip + il / 16; + + const uint y_idx = i * QUANT_K + 128 * ip + il; + + const uint ql_idx = 32 * ip + il; + const uint8_t qs = data_a[i].qs[32 * ip + il]; + + FLOAT_TYPE dall = FLOAT_TYPE(data_a[i].d.x); + FLOAT_TYPE dmin = FLOAT_TYPE(data_a[i].d.y); + data_b[y_idx + 0] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+0] & 0xF) * ((qs >> 0) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+0] >> 4)); + data_b[y_idx + 32] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+2] & 0xF) * ((qs >> 2) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+2] >> 4)); + data_b[y_idx + 64] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+4] & 0xF) * ((qs >> 4) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+4] >> 4)); + data_b[y_idx + 96] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+6] & 0xF) * ((qs >> 6) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+6] >> 4)); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp new file mode 100644 index 000000000..c17dd0d99 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp @@ -0,0 +1,42 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { + const uint i = uint(gl_WorkGroupID.x * 256 + wgy); + if (i >= p.M * p.K / QUANT_K) { + return; + } + + const uint r = gl_LocalInvocationID.x / 4; + const uint tid = r / 2; + const uint is0 = r % 2; + const uint l0 = 16 * is0 + 4 * (gl_LocalInvocationID.x % 4); + const uint n = tid / 4; + const uint j = tid - 4*n; + + const uint8_t m = uint8_t(1 << (4*n + j)); + const uint is = 8*n + 2*j + is0; + const uint shift = 2*j; + + const int8_t us = int8_t(is < 4 ? (data_a[i].scales[is-0] & 0xF) | (((data_a[i].scales[is+8] >> 0) & 3) << 4) : + is < 8 ? (data_a[i].scales[is-0] & 0xF) | (((data_a[i].scales[is+4] >> 2) & 3) << 4) : + is < 12 ? (data_a[i].scales[is-8] >> 4) | (((data_a[i].scales[is+0] >> 4) & 3) << 4) : + (data_a[i].scales[is-8] >> 4) | (((data_a[i].scales[is-4] >> 6) & 3) << 4)); + const FLOAT_TYPE d_all = FLOAT_TYPE(data_a[i].d); + const FLOAT_TYPE dl = d_all * FLOAT_TYPE(us - 32); + + const uint y_idx = i * QUANT_K + 128 * n + 32 * j; + const uint qs_idx = 32*n; + + for (uint l = l0; l < l0 + 4; ++l) { + data_b[y_idx + l] = D_TYPE(dl * FLOAT_TYPE(int8_t((data_a[i].qs[qs_idx + l] >> shift) & 3) - (((data_a[i].hmask[l] & m) != 0) ? 0 : 4))); + } + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp new file mode 100644 index 000000000..408185327 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp @@ -0,0 +1,30 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_q4_0 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint q_idx = 8*il; + const uint b_idx = 1024*i + 32*ir + q_idx; + + const float d = float(data_a[ib].d); + + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_b[b_idx + l + 0] = D_TYPE(d * ((data_a[ib].qs[q_idx + l] & 0xF) - 8.0f)); + data_b[b_idx + l + 16] = D_TYPE(d * ((data_a[ib].qs[q_idx + l] >> 4) - 8.0f)); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp new file mode 100644 index 000000000..2f27eee68 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp @@ -0,0 +1,32 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_q4_1 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint b_idx = 1024*i + 32*ir + 8*il; + + const float d = float(data_a[ib].d); + const float m = float(data_a[ib].m); + + const uint q_idx = 8*il; + + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_b[b_idx + l + 0] = D_TYPE(d * (data_a[ib].qs[q_idx + l] & 0xF) + m); + data_b[b_idx + l + 16] = D_TYPE(d * (data_a[ib].qs[q_idx + l] >> 4) + m); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp new file mode 100644 index 000000000..987f113a3 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp @@ -0,0 +1,68 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { + const uint ib = gl_WorkGroupID.x * 256 + wgy; + if (ib >= p.M * p.K / QUANT_K) { + return; + } + + const uint tid = gl_LocalInvocationID.x; + const uint il = tid / 8; + const uint ir = tid % 8; + const uint is = 2 * il; + const uint n = 4; + + const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib].d.x); + const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib].d.y); + + const uint y_idx = ib * QUANT_K + 64 * il + n * ir; + const uint qs_idx = 32*il + n * ir; + + uint scidx0 = (is < 4) ? is : (is + 4); + uint scidx1 = (is < 4) ? is : (is - 4); + uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint scidxshift1 = (is < 4) ? 0 : 2; + uint mbidx0 = is + 4; + uint mbidx1 = (is < 4) ? is + 4 : is; + uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; + uint mbidxshift0 = (is < 4) ? 0 : 4; + uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint mbidxshift1 = (is < 4) ? 0 : 2; + + uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const FLOAT_TYPE d1 = dall * sc; + const FLOAT_TYPE m1 = dmin * mbyte; + + scidx0 = (is < 4) ? is + 1 : (is + 5); + scidx1 = (is < 4) ? is + 1 : (is - 3); + scidxmask1 = (is < 4) ? 0x30 : 0xC0; + scidxshift1 = (is < 4) ? 0 : 2; + mbidx0 = is + 5; + mbidx1 = (is < 4) ? is + 5 : is + 1; + mbidxmask0 = (is < 4) ? 0xF : 0xF0; + mbidxshift0 = (is < 4) ? 0 : 4; + mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + mbidxshift1 = (is < 4) ? 0 : 2; + + sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const FLOAT_TYPE d2 = dall * sc; + const FLOAT_TYPE m2 = dmin * mbyte; + + [[unroll]] for (uint l = 0; l < n; ++l) { + data_b[y_idx + l ] = D_TYPE(d1 * FLOAT_TYPE(data_a[ib].qs[qs_idx + l] & 0xF) - m1); + data_b[y_idx + l + 32] = D_TYPE(d2 * FLOAT_TYPE(data_a[ib].qs[qs_idx + l] >> 4) - m2); + } + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp new file mode 100644 index 000000000..b20b80529 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp @@ -0,0 +1,34 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_q5_0 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint b_idx = 1024*i + 32*ir + 8*il; + + const float d = float(data_a[ib].d); + const uint qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0]; + + const uint q_idx = 8*il; + + [[unroll]] for (uint l = 0; l < 8; ++l) { + const uint iqs = q_idx + l; + const uint vui = uint(data_a[ib].qs[iqs]); + data_b[b_idx + l + 0] = D_TYPE(d * (((vui & 0xF) | (((qh >> iqs) << 4) & 0x10)) - 16.0f)); + data_b[b_idx + l + 16] = D_TYPE(d * (((vui >> 4) | ((qh >> (iqs + 12)) & 0x10)) - 16.0f)); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp new file mode 100644 index 000000000..dc59fe3b7 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp @@ -0,0 +1,35 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_q5_1 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint b_idx = 1024*i + 32*ir + 8*il; + + const float d = float(data_a[ib].d); + const float m = float(data_a[ib].m); + const uint qh = data_a[ib].qh; + + const uint q_idx = 8*il; + + [[unroll]] for (uint l = 0; l < 8; ++l) { + const uint iqs = q_idx + l; + const uint vui = uint(data_a[ib].qs[iqs]); + data_b[b_idx + l + 0] = D_TYPE(d * (((vui & 0xF) | (((qh >> iqs) << 4) & 0x10))) + m); + data_b[b_idx + l + 16] = D_TYPE(d * (((vui >> 4) | ((qh >> (iqs + 12)) & 0x10))) + m); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp new file mode 100644 index 000000000..6db5403b6 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp @@ -0,0 +1,70 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { + const uint ib = gl_WorkGroupID.x * 256 + wgy; + if (ib >= p.M * p.K / QUANT_K) { + return; + } + + const uint tid = gl_LocalInvocationID.x; + const uint il = tid / 16; + const uint ir = tid % 16; + const uint is = 2 * il; + + const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib].d.x); + const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib].d.y); + + const uint y_idx = ib * QUANT_K + 64 * il + 2 * ir; + const uint qs_idx = 32*il + 2 * ir; + const uint qh_idx = 2 * ir; + + uint scidx0 = (is < 4) ? is : (is + 4); + uint scidx1 = (is < 4) ? is : (is - 4); + uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint scidxshift1 = (is < 4) ? 0 : 2; + uint mbidx0 = is + 4; + uint mbidx1 = (is < 4) ? is + 4 : is; + uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; + uint mbidxshift0 = (is < 4) ? 0 : 4; + uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint mbidxshift1 = (is < 4) ? 0 : 2; + + uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const FLOAT_TYPE d1 = dall * sc; + const FLOAT_TYPE m1 = dmin * mbyte; + + scidx0 = (is < 4) ? is + 1 : (is + 5); + scidx1 = (is < 4) ? is + 1 : (is - 3); + scidxmask1 = (is < 4) ? 0x30 : 0xC0; + scidxshift1 = (is < 4) ? 0 : 2; + mbidx0 = is + 5; + mbidx1 = (is < 4) ? is + 5 : is + 1; + mbidxmask0 = (is < 4) ? 0xF : 0xF0; + mbidxshift0 = (is < 4) ? 0 : 4; + mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + mbidxshift1 = (is < 4) ? 0 : 2; + + sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const FLOAT_TYPE d2 = dall * sc; + const FLOAT_TYPE m2 = dmin * mbyte; + + const uint8_t hm1 = uint8_t(1 << (2 * il )); + const uint8_t hm2 = uint8_t(1 << (2 * il + 1)); + data_b[y_idx ] = D_TYPE(d1 * FLOAT_TYPE((data_a[ib].qs[qs_idx ] & 0xF) + (((data_a[ib].qh[qh_idx ] & hm1) != 0) ? 16 : 0)) - m1); + data_b[y_idx + 1] = D_TYPE(d1 * FLOAT_TYPE((data_a[ib].qs[qs_idx + 1] & 0xF) + (((data_a[ib].qh[qh_idx + 1] & hm1) != 0) ? 16 : 0)) - m1); + data_b[y_idx + 32] = D_TYPE(d2 * FLOAT_TYPE((data_a[ib].qs[qs_idx ] >> 4) + (((data_a[ib].qh[qh_idx ] & hm2) != 0) ? 16 : 0)) - m2); + data_b[y_idx + 33] = D_TYPE(d2 * FLOAT_TYPE((data_a[ib].qs[qs_idx + 1] >> 4) + (((data_a[ib].qh[qh_idx + 1] & hm2) != 0) ? 16 : 0)) - m2); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp new file mode 100644 index 000000000..0b9131755 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp @@ -0,0 +1,33 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { + const uint i = gl_WorkGroupID.x * 256 + wgy; + if (i >= p.M * p.K / QUANT_K) { + return; + } + const uint tid = gl_LocalInvocationID.x; + const uint ip = tid / 32; + const uint il = tid - 32 * ip; + const uint is = 8 * ip + il / 16; + + const uint y_idx = i * QUANT_K + 128 * ip + il; + + const uint ql_idx = 64 * ip + il; + const uint8_t qh = data_a[i].qh[32 * ip + il]; + + const FLOAT_TYPE d = FLOAT_TYPE(data_a[i].d); + + data_b[y_idx + 0] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 0] * (int8_t((data_a[i].ql[ql_idx + 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32))); + data_b[y_idx + 32] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 2] * (int8_t((data_a[i].ql[ql_idx + 32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32))); + data_b[y_idx + 64] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 4] * (int8_t((data_a[i].ql[ql_idx + 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32))); + data_b[y_idx + 96] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 6] * (int8_t((data_a[i].ql[ql_idx + 32] >> 4) | (((qh >> 6) & 3) << 4)) - 32))); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp new file mode 100644 index 000000000..bd1344a88 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp @@ -0,0 +1,31 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_q8_0 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint b_idx = 1024*i + 32*ir + 16*il; + + const float d = float(data_a[ib].d); + + const uint q_idx = 16*il; + + [[unroll]] for (uint l = 0; l < 16; l += 2) { + data_b[b_idx + l ] = D_TYPE(d * data_a[ib].qs[q_idx + l ]); + data_b[b_idx + l + 1] = D_TYPE(d * data_a[ib].qs[q_idx + l + 1]); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp new file mode 100644 index 000000000..4e68742b5 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp @@ -0,0 +1,34 @@ +#version 450 + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_control_flow_attributes : enable + +layout (push_constant) uniform parameter +{ + uint ncols; + uint rows_per_channel; + uint n_past; +} p; + +#include "types.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint col = gl_GlobalInvocationID.y; + const uint row = gl_GlobalInvocationID.x; + + if (col >= p.ncols) { + return; + } + + const uint i = row*p.ncols + col; + if (col > p.n_past + row % p.rows_per_channel) { + data_d[i] = D_TYPE(uintBitsToFloat(0xFF800000)); + } else { + data_d[i] = D_TYPE(data_a[i]); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp new file mode 100644 index 000000000..9fb69c6c1 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp @@ -0,0 +1,27 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" + +const uint num_threads = 256; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + uint idx = get_idx(); + + // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation + const uint num_iter = 2; + + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (idx >= p.ne) { + continue; + } + uint i00, i01, i02, i03; + get_indices(idx, i00, i01, i02, i03); + + data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) / FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); + + idx += num_threads; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp new file mode 100644 index 000000000..c5be8131b --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp @@ -0,0 +1,289 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require + +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +#extension GL_KHR_memory_scope_semantics : enable +#extension GL_KHR_cooperative_matrix : enable +#extension GL_NV_cooperative_matrix2 : enable +#extension GL_EXT_buffer_reference : enable +#extension GL_KHR_shader_subgroup_ballot : enable +#extension GL_KHR_shader_subgroup_vote : enable +#extension GL_EXT_null_initializer : enable + +#include "types.comp" +#include "dequant_funcs_cm2.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (constant_id = 1) const uint32_t Br = 32; +layout (constant_id = 2) const uint32_t Bc = 32; +layout (constant_id = 3) const uint32_t D = 32; +layout (constant_id = 4) const uint32_t Clamp = gl_CooperativeMatrixClampModeConstantNV; + +layout (push_constant) uniform parameter { + uint32_t N; + uint32_t KV; + + uint32_t ne1; + uint32_t ne2; + uint32_t ne3; + + uint32_t neq2; + uint32_t neq3; + uint32_t nek2; + uint32_t nek3; + uint32_t nev2; + uint32_t nev3; + uint32_t nem1; + + uint32_t nb02; + uint32_t nb03; + uint32_t nb12; + uint32_t nb13; + uint32_t nb22; + uint32_t nb23; + uint32_t nb31; + + float scale; + float max_bias; + float logit_softcap; + + uint32_t mask; + uint32_t n_head_log2; + float m0; + float m1; +} p; + +layout (binding = 0) readonly buffer Q {uint8_t data_q[];}; +layout (binding = 1) readonly buffer K {uint8_t data_k[];}; +layout (binding = 2) readonly buffer V {uint8_t data_v[];}; +layout (binding = 3) readonly buffer M {uint8_t data_m[];}; +layout (binding = 4) writeonly buffer O {D_TYPE data_o[];}; + +#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b)) + +ACC_TYPE maxReduce(const in ACC_TYPE x, const in ACC_TYPE y) { + return max(x, y); +} + +ACC_TYPE smearReduce(const in ACC_TYPE x, const in ACC_TYPE y) { + return x; +} + +// Replace matrix elements >= numRows or numCols with 'replace' +ACC_TYPE replacePadding(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem, const in ACC_TYPE replace, const in uint32_t numRows, const in uint32_t numCols) { + if (row >= numRows || col >= numCols) { + return replace; + } + return elem; +} + +ACC_TYPE Exp(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem) +{ + return exp(elem); +} + +ACC_TYPE Max(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem0, const in ACC_TYPE elem1) +{ + return max(elem0, elem1); +} + +#if defined(BLOCK_SIZE) +#define DECODEFUNC , DEQUANTFUNC +#else +#define DECODEFUNC +#endif + +void main() { +#if defined(DATA_A_IQ4_NL) + init_iq4nl_shmem(); +#endif + + const uint32_t N = p.N; + const uint32_t KV = p.KV; + + const uint32_t Tr = CEIL_DIV(N, Br); + const uint32_t Tc = CEIL_DIV(KV, Bc); + + const uint32_t i = gl_WorkGroupID.x; + + const uint32_t iq2 = gl_WorkGroupID.y; + const uint32_t iq3 = gl_WorkGroupID.z; + + // broadcast factors + const uint32_t rk2 = p.neq2/p.nek2; + const uint32_t rk3 = p.neq3/p.nek3; + + const uint32_t rv2 = p.neq2/p.nev2; + const uint32_t rv3 = p.neq3/p.nev3; + + // k indices + const uint32_t ik3 = iq3 / rk3; + const uint32_t ik2 = iq2 / rk2; + + // v indices + const uint32_t iv3 = iq3 / rv3; + const uint32_t iv2 = iq2 / rv2; + + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutQ = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutNV<2, Clamp> tensorLayoutK = createTensorLayoutNV(2, Clamp); + tensorLayoutNV<2, Clamp> tensorLayoutV = createTensorLayoutNV(2, Clamp); + + tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0); + +#if defined(BLOCK_SIZE) + tensorLayoutK = setTensorLayoutBlockSizeNV(tensorLayoutK, 1, BLOCK_SIZE); + tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, BLOCK_SIZE); +#endif + + tensorLayoutQ = setTensorLayoutDimensionNV(tensorLayoutQ, N, D); + tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, KV, D); + tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, KV, D); + + coopmat Q; + coopmat Qf16; + + uint32_t q_offset = iq2*p.nb02+iq3*p.nb03; + coopMatLoadTensorNV(Q, data_q, q_offset, sliceTensorLayoutNV(tensorLayoutQ, i * Br, Br, 0, D)); + + Qf16 = coopmat(Q); + Qf16 *= float16_t(p.scale); + + coopmat O = coopmat(0); + + coopmat L, M; + + L = coopmat(0); + M = coopmat(-1.0/0.0); + + ACC_TYPE slope = ACC_TYPE(1.0); + + // ALiBi + if (p.max_bias > 0.0f) { + const uint32_t h = iq2; + + const ACC_TYPE base = ACC_TYPE(h < p.n_head_log2 ? p.m0 : p.m1); + const int exph = int(h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1); + + slope = pow(base, ACC_TYPE(exph)); + } + + [[dont_unroll]] + for (uint32_t j = 0; j < Tc; ++j) { + + coopmat S = coopmat(0); + + coopmat K_T; + + uint32_t k_offset = ik2*p.nb12 + ik3*p.nb13; + coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, D), tensorViewTranspose DECODEFUNC); + S = coopMatMulAdd(Qf16, K_T, S); + + if (p.logit_softcap != 0.0f) { + [[unroll]] + for (int k = 0; k < S.length(); ++k) { + S[k] = ACC_TYPE(p.logit_softcap)*tanh(S[k]); + } + } + + if (p.mask != 0) { + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutM = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutM = setTensorLayoutDimensionNV(tensorLayoutM, p.nem1, KV); + + coopmat mv; + + coopMatLoadTensorNV(mv, data_m, 0, sliceTensorLayoutNV(tensorLayoutM, i * Br, Br, j * Bc, Bc)); + + S += slope*coopmat(mv); + } + + // Clear padding elements to -inf, so they don't contribute to rowmax + if (Clamp != 0 && + ((j + 1) * Bc > KV || + (i + 1) * Br > N)) { + + uint R = ((i + 1) * Br > N) ? (N % Br) : Br; + uint C = ((j + 1) * Bc > KV) ? (KV % Bc) : Bc; + + coopMatPerElementNV(S, S, replacePadding, ACC_TYPE(-1.0/0.0), R, C); + } + + coopmat rowmax, P, rowsum, eM; + + coopMatReduceNV(rowmax, S, gl_CooperativeMatrixReduceRowNV, maxReduce); + + coopmat Mold = M; + + // M = max(rowmax, Mold) + // P = e^(S - M) + // eM = e^(Mold - M) + coopMatPerElementNV(M, rowmax, Max, Mold); + coopMatPerElementNV(P, S - M, Exp); + coopMatPerElementNV(eM, Mold - M, Exp); + + // Clear padding elements to 0, so they don't contribute to rowsum + if (Clamp != 0 && + ((j + 1) * Bc > KV || + (i + 1) * Br > N)) { + + uint R = ((i + 1) * Br > N) ? (N % Br) : Br; + uint C = ((j + 1) * Bc > KV) ? (KV % Bc) : Bc; + + coopMatPerElementNV(P, P, replacePadding, ACC_TYPE(0.0), R, C); + } + + coopmat P_A = coopmat(P); + + // compute rowsum by multiplying by matrix of all ones. + coopmat One = coopmat(1.0); + + rowsum = coopmat(0.0); + rowsum = coopMatMulAdd(P_A, One, rowsum); + + coopmat V; + uint32_t v_offset = iv2*p.nb22 + iv3*p.nb23; + coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, D) DECODEFUNC); + + L = eM*L + rowsum; + + // This is the "diagonal" matrix in the paper, but since we do componentwise + // multiply rather than matrix multiply it has the diagonal element smeared + // across the row + coopmat eMdiag; + + // resize eM by using smear/reduce + coopMatReduceNV(eMdiag, eM, gl_CooperativeMatrixReduceRowNV, smearReduce); + + O = eMdiag * O; + + O = coopMatMulAdd(P_A, V, O); + } + + coopmat Ldiag; + + // resize L by using smear/reduce + coopMatReduceNV(Ldiag, L, gl_CooperativeMatrixReduceRowNV, smearReduce); + + [[unroll]] + for (int k = 0; k < Ldiag.length(); ++k) { + Ldiag[k] = ACC_TYPE(1.0) / Ldiag[k]; + } + + O = Ldiag*O; + + tensorLayoutNV<3, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(3, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.ne2, p.ne1, D); + + // permute dimensions + tensorViewNV<3, false, 1, 0, 2> tensorViewPermute = createTensorViewNV(3, false, 1, 0, 2); + uint32_t o_offset = iq3*p.ne2*p.ne1; + + coopmat O_D = coopmat(O); + coopMatStoreTensorNV(O_D, data_o, o_offset, sliceTensorLayoutNV(tensorLayoutD, i * Br, Br, iq2, 1, 0, D), tensorViewPermute); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp new file mode 100644 index 000000000..4cc7a68ca --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp @@ -0,0 +1,25 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const float GELU_COEF_A = 0.044715f; + const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + + const float xi = float(data_a[i]); + const float val = SQRT_2_OVER_PI*xi*(1.0f + GELU_COEF_A*xi*xi); + data_d[i] = D_TYPE(0.5f*xi*(2.0f - 2.0f / (exp(2 * val) + 1))); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp new file mode 100644 index 000000000..e6e6fcfd2 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp @@ -0,0 +1,23 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const float GELU_QUICK_COEF = -1.702f; + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + + const float x = float(data_a[i]); + data_d[i] = D_TYPE(x * (1.0f / (1.0f + exp(GELU_QUICK_COEF * x)))); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp new file mode 100644 index 000000000..062e2a4cd --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp @@ -0,0 +1,64 @@ +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_control_flow_attributes : require + +layout (push_constant) uniform parameter +{ + uint ne; + uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03; + uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13; + uint ne20; uint ne21; uint ne22; uint ne23; uint nb20; uint nb21; uint nb22; uint nb23; + uint misalign_offsets; + float param1; float param2; int param3; +} p; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; + +// true if src0/src1 are the same shape and the indices can be reused without additional modulus +layout(constant_id = 0) const bool norepeat = false; + +uint get_idx() { + return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; +} + +uint get_aoffset() { return p.misalign_offsets >> 16; } +uint get_boffset() { return (p.misalign_offsets >> 8) & 0xFF; } +uint get_doffset() { return p.misalign_offsets & 0xFF; } + +// mod and div are expensive and coordinates/dimensions are often power of 2 or equal to 1 +uint fastmod(uint a, uint b) { + if ((b & (b-1)) == 0) { + return a & (b-1); + } + return a % b; +} + +uint fastdiv(uint a, uint b) { + return (a < b) ? 0 : (a / b); +} + +void get_indices(uint idx, out uint i00, out uint i01, out uint i02, out uint i03) { + i03 = fastdiv(idx, (p.ne02*p.ne01*p.ne00)); + const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00; + i02 = fastdiv((idx - i03_offset), (p.ne01*p.ne00)); + const uint i02_offset = i02*p.ne01*p.ne00; + i01 = (idx - i03_offset - i02_offset) / p.ne00; + i00 = idx - i03_offset - i02_offset - i01*p.ne00; +} + +uint src0_idx(uint i00, uint i01, uint i02, uint i03) { + return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00; +} + +uint src1_idx(uint i00, uint i01, uint i02, uint i03) { + if (norepeat) { + return i03*p.nb13 + i02*p.nb12 + i01*p.nb11 + i00*p.nb10; + } else { + return fastmod(i03, p.ne13)*p.nb13 + fastmod(i02, p.ne12)*p.nb12 + fastmod(i01, p.ne11)*p.nb11 + fastmod(i00, p.ne10)*p.nb10; + } +} + +uint dst_idx(uint i00, uint i01, uint i02, uint i03) { + return i03*p.nb23 + i02*p.nb22 + i01*p.nb21 + i00*p.nb20; +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp new file mode 100644 index 000000000..66e46ae67 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp @@ -0,0 +1,9 @@ +#extension GL_EXT_shader_16bit_storage : require + +layout (push_constant) uniform parameter +{ + uint KX; + uint KY; + float param1; + float param2; +} p; diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp new file mode 100644 index 000000000..68d1bc9f1 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp @@ -0,0 +1,56 @@ +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_control_flow_attributes : require + +layout (push_constant) uniform parameter +{ + uint ne; + uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03; + uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13; + uint misalign_offsets; + float param1; float param2; + + uint ne0_012mp; uint ne0_012L; + uint ne0_01mp; uint ne0_01L; + uint ne0_0mp; uint ne0_0L; + uint ne1_012mp; uint ne1_012L; + uint ne1_01mp; uint ne1_01L; + uint ne1_0mp; uint ne1_0L; +} p; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +uint get_idx() { + return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; +} + +uint get_aoffset() { return p.misalign_offsets >> 16; } +uint get_doffset() { return p.misalign_offsets & 0xFFFF; } + +// see init_fastdiv_values in ggml-vulkan.cpp +uint fastdiv(uint n, uint mp, uint L) { + uint msbs, lsbs; + // msbs = mulhi(n, mp) + umulExtended(n, mp, msbs, lsbs); + return (msbs + n) >> L; +} + +uint src0_idx(uint idx) { + const uint i03 = fastdiv(idx, p.ne0_012mp, p.ne0_012L); + const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00; + const uint i02 = fastdiv(idx - i03_offset, p.ne0_01mp, p.ne0_01L); + const uint i02_offset = i02*p.ne01*p.ne00; + const uint i01 = fastdiv(idx - i03_offset - i02_offset, p.ne0_0mp, p.ne0_0L); + const uint i00 = idx - i03_offset - i02_offset - i01*p.ne00; + return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00; +} + +uint dst_idx(uint idx) { + const uint i13 = fastdiv(idx, p.ne1_012mp, p.ne1_012L); + const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10; + const uint i12 = fastdiv(idx - i13_offset, p.ne1_01mp, p.ne1_01L); + const uint i12_offset = i12*p.ne11*p.ne10; + const uint i11 = fastdiv(idx - i13_offset - i12_offset, p.ne1_0mp, p.ne1_0L); + const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10; + return i13*p.nb13 + i12*p.nb12 + i11*p.nb11 + i10*p.nb10; +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp new file mode 100644 index 000000000..e877ed779 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp @@ -0,0 +1,28 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint i00 = gl_GlobalInvocationID.x; + const uint i10 = gl_GlobalInvocationID.y; + const uint i11 = (gl_GlobalInvocationID.z)/p.ne12; + const uint i12 = (gl_GlobalInvocationID.z)%p.ne12; + + if (i00 >= p.ne00) { + return; + } + + const uint i01 = data_b[get_boffset() + i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; + + const uint a_offset = get_aoffset() + i01*p.nb01 + i11*p.nb02 + i12*p.nb03; + const uint d_offset = get_doffset() + i10*p.nb21 + i11*p.nb22 + i12*p.nb23; + +#ifndef OPTIMIZATION_ERROR_WORKAROUND + data_d[d_offset + i00] = D_TYPE(data_a[a_offset + i00]); +#else + data_d[d_offset + i00] = data_a[a_offset + i00]; +#endif +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp new file mode 100644 index 000000000..1426fde65 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp @@ -0,0 +1,39 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" +#include "dequant_funcs.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint i00 = (gl_GlobalInvocationID.x)*2; + const uint i10 = gl_GlobalInvocationID.y; + const uint i11 = (gl_GlobalInvocationID.z)/p.ne12; + const uint i12 = (gl_GlobalInvocationID.z)%p.ne12; + +#if defined(DATA_A_IQ4_NL) + init_iq4nl_shmem(); +#endif + + if (i00 >= p.ne00) { + return; + } + + const uint i01 = data_b[i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; + + const uint a_offset = i01*p.nb01 + i11*p.nb02 + i12*p.nb03; + const uint d_offset = i10*p.nb21 + i11*p.nb22 + i12*p.nb23; + + const uint ib = a_offset + i00/QUANT_K; // block index + const uint iqs = (i00%QUANT_K)/QUANT_R; // quant index + const uint iybs = i00 - i00%QUANT_K; // dst block start index + const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2; + + vec2 v = dequantize(ib, iqs, 0); + const vec2 dm = get_dm(ib, 0); + v = v * dm.x + dm.y; + + data_d[d_offset + iybs + iqs ] = D_TYPE(v.x); + data_d[d_offset + iybs + iqs + y_offset] = D_TYPE(v.y); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp new file mode 100644 index 000000000..b6a0d5645 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp @@ -0,0 +1,66 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable +#define BLOCK_SIZE 512 + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +shared float tmp[BLOCK_SIZE]; + +void main() { + const uint group_size = p.KX; + const float eps = p.param1; + + const uint tid = gl_LocalInvocationID.x; + const uint start = gl_WorkGroupID.x * group_size + tid; + const uint end = (gl_WorkGroupID.x + 1) * group_size; + + tmp[tid] = 0.0f; + + // Calculate mean + [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { + tmp[tid] += float(data_a[col]); + } + + // tmp up partial tmps and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] += tmp[tid + s]; + } + barrier(); + } + + const float mean = tmp[0] / group_size; + barrier(); + tmp[tid] = 0.0f; + + // Calculate variance + [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { + const float xi = float(data_a[col]) - mean; + data_d[col] = D_TYPE(xi); + tmp[tid] += xi * xi; + } + + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] += tmp[tid + s]; + } + barrier(); + } + + const float variance = tmp[0] / group_size; + const float scale = inversesqrt(variance + eps); + + [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { + data_d[col] *= D_TYPE(scale); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp new file mode 100644 index 000000000..122b1e93f --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp @@ -0,0 +1,87 @@ +#version 450 + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_spirv_intrinsics: enable +#extension GL_EXT_control_flow_attributes : require + +#if RTE16 +spirv_execution_mode(capabilities = [4467], 4462, 16); // RoundingModeRTE, 16 bits +#endif + +layout (push_constant) uniform parameter +{ + uint batch_offset; uint offset_delta; + uint IC; + uint IW; uint IH; + uint OW; uint OH; + uint KW; uint KH; + uint pelements; + uint CHW; + int s0; int s1; + int p0; int p1; + int d0; int d1; +} p; + +#include "types.comp" + +layout(constant_id = 0) const uint BLOCK_SIZE = 32; + +const uint NUM_ITER = 512 / BLOCK_SIZE; + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint gidx = gl_GlobalInvocationID.x; + + const uint oh = gl_GlobalInvocationID.y; + const uint batch = gl_GlobalInvocationID.z / p.IC; + const uint ic = gl_GlobalInvocationID.z % p.IC; + + A_TYPE values[NUM_ITER]; + uint offset_dst[NUM_ITER]; + [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { + values[idx] = A_TYPE(0); + } + + [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { + + const uint i = gidx * NUM_ITER + idx; + + const uint ksize = p.OW * (p.KH > 1 ? p.KW : 1); + const uint kx = i / ksize; + const uint kd = kx * ksize; + const uint ky = (i - kd) / p.OW; + const uint ix = i % p.OW; + + const uint iiw = ix * p.s0 + kx * p.d0 - p.p0; + const uint iih = oh * p.s1 + ky * p.d1 - p.p1; + + offset_dst[idx] = + ((batch * p.OH + oh) * p.OW + ix) * p.CHW + + (ic * (p.KW * p.KH) + ky * p.KW + kx); + + if (i >= p.pelements) { + continue; + } + + if (iih < p.IH && iiw < p.IW) { + const uint offset_src = ic * p.offset_delta + batch * p.batch_offset; + values[idx] = data_a[offset_src + iih * p.IW + iiw]; + } + } + + [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { + + const uint i = gidx * NUM_ITER + idx; + + if (i >= p.pelements) { + continue; + } + + data_d[offset_dst[idx]] = D_TYPE(values[idx]); + } + +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp new file mode 100644 index 000000000..d90a99aea --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp @@ -0,0 +1,22 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + + const float val = float(data_a[i]); + data_d[i] = D_TYPE(max(val, 0.0f) + min(val, 0.0f) * p.param1); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp new file mode 100644 index 000000000..43de19df8 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp @@ -0,0 +1,27 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" + +const uint num_threads = 256; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + uint idx = get_idx(); + + // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation + const uint num_iter = 2; + + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (idx >= p.ne) { + continue; + } + uint i00, i01, i02, i03; + get_indices(idx, i00, i01, i02, i03); + + data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) * FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); + + idx += num_threads; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp new file mode 100644 index 000000000..4c64fd47a --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp @@ -0,0 +1,48 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {float data_a[];}; +layout (binding = 0) readonly buffer A4 {vec4 data_a4[];}; +layout (binding = 1) writeonly buffer D {float data_d[];}; +layout (binding = 1) writeonly buffer D4 {vec4 data_d4[];}; + +layout (push_constant) uniform parameter { + uint ne; + uint k_num; +} p; + +void main() { + // Each invocation handles four consecutive components + const uint idx = gl_GlobalInvocationID.x * 4; + + if (idx >= p.ne) { + return; + } + + // Check if all four components are in bounds and aligned, + // then use vector loads + if (idx + 3 < p.ne && (p.ne % 4) == 0) { + vec4 result = vec4(0.0f); + + [[unroll]] for (uint i = 0; i < p.k_num; i++) { + result += data_a4[(i * p.ne + idx) / 4]; + } + + data_d4[idx / 4] = result; + } else { + [[unroll]] for (uint j = 0; j < 4; ++j) { + if (idx + j < p.ne) { + float result = 0.0f; + + [[unroll]] for (uint i = 0; i < p.k_num; i++) { + result += data_a[i * p.ne + idx + j]; + } + + data_d[idx + j] = result; + } + } + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp new file mode 100644 index 000000000..24875cdcf --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp @@ -0,0 +1,152 @@ +#version 450 + +#ifdef FLOAT16 +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#endif +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +#if !defined(DATA_A_F32) && !defined(DATA_A_F16) +#define K_PER_ITER 8 +#else +#define K_PER_ITER 2 +#endif + + +uint a_offset, b_offset, d_offset, y_offset; + +void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint tid, const uint i, bool lastiter) +{ + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + const uint col = i*BLOCK_SIZE + K_PER_ITER*tid; + const uint iqs = (col%QUANT_K)/QUANT_R; // quant index + const uint iybs = col - col%QUANT_K; // y block start index + +#if K_PER_ITER == 8 +#if QUANT_R == 2 + const B_TYPE_VEC4 bv02 = data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]; + const B_TYPE_VEC4 bv13 = data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs + y_offset) / 4]; + const vec4 bv0 = vec4(bv02.x, bv13.x, bv02.y, bv13.y); + const vec4 bv1 = vec4(bv02.z, bv13.z, bv02.w, bv13.w); +#else + const vec4 bv0 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]); + const vec4 bv1 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4 + 1]); +#endif +#else + // Check if the second of the pair of elements is OOB, and don't fetch B or + // accumulate it. We still fetch a pair of elements for A, which is fine for + // quantized formats since they'll be within the same block. We should + // probably skip fetching the second element for F16/F32, but as of now we + // still do. + const bool OOB = lastiter && (iybs + iqs + y_offset >= p.ncols); + + FLOAT_TYPE b0 = 0, b1 = 0; + b0 = FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs]); + if (!OOB) { + b1 = FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs + y_offset]); + } +#endif + uint ibi = first_row*p.ncols; + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib = (ibi + col)/QUANT_K; // block index + ibi += p.ncols; + +#if K_PER_ITER == 8 + vec4 v = dequantize4(ib, iqs, a_offset); + vec4 v2 = dequantize4(ib, iqs+(4/QUANT_R), a_offset); + + const vec2 dm = get_dm(ib, a_offset); + if (dm.y != 0) { // quant has min component + v = v * dm.x + dm.y; + v2 = v2 * dm.x + dm.y; + } + + // matrix multiplication + FLOAT_TYPE rowtmp = dot(bv0, v); + rowtmp += dot(bv1, v2); + + if (dm.y == 0) + rowtmp *= dm.x; + + temp[j][n] += rowtmp; +#else + const vec2 v = dequantize(ib, iqs, a_offset); + + // matrix multiplication + temp[j][n] = fma(FLOAT_TYPE(v.x), b0, temp[j][n]); + if (!OOB) { + temp[j][n] = fma(FLOAT_TYPE(v.y), b1, temp[j][n]); + } +#endif + } + } +} + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + const uint tid = gl_LocalInvocationID.x; + + get_offsets(a_offset, b_offset, d_offset); + a_offset /= QUANT_K; + + y_offset = QUANT_R == 1 ? 1 : QUANT_K/2; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + uint num_iters = p.ncols / (K_PER_ITER * BLOCK_SIZE); + if (num_iters * K_PER_ITER * BLOCK_SIZE + K_PER_ITER*tid < p.ncols) { + num_iters++; + } + int unroll_count = 4; + uint unrolled_iters = num_iters & ~(unroll_count - 1); + + uint i = 0; + while (i < unrolled_iters) { + // Manually partially unroll the loop + [[unroll]] for (uint k = 0; k < unroll_count; ++k) { + iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false); + i++; + } + } + unroll_count = 2; + unrolled_iters = num_iters & ~(unroll_count - 1); + while (i < unrolled_iters) { + // Manually partially unroll the loop + [[unroll]] for (uint k = 0; k < unroll_count; ++k) { + iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false); + i++; + } + } + while (i < num_iters) { + iter(temp, first_row, num_rows, tid, i*K_PER_ITER, true); + i++; + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + +#if defined(DATA_A_IQ4_NL) + init_iq4nl_shmem(); +#endif + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp new file mode 100644 index 000000000..903753c7e --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp @@ -0,0 +1,118 @@ +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_8bit_storage : require + +#ifdef MUL_MAT_ID +#define EXPERT_COUNT 8 +#endif + +#include "types.comp" + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 1) readonly buffer BV2 {B_TYPE_VEC2 data_b_v2[];}; +layout (binding = 1) readonly buffer BV4 {B_TYPE_VEC4 data_b_v4[];}; + +layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; +#ifdef MUL_MAT_ID +layout (binding = 3) readonly buffer IDS {int data_ids[];}; +#endif + +#include "dequant_funcs.comp" + +layout (push_constant) uniform parameter +{ + uint ncols; + uint stride_a; + uint stride_b; + uint stride_d; + + uint batch_stride_a; + uint batch_stride_b; + uint batch_stride_d; + +#ifdef MUL_MAT_ID + uint nei0; + uint ne11; +#else + uint ne02; + uint ne12; + uint broadcast2; + uint broadcast3; +#endif +} p; + +void get_offsets(out uint a_offset, out uint b_offset, out uint d_offset) { +#ifdef MUL_MAT_ID + const uint expert_idx = gl_GlobalInvocationID.y; +#else + const uint batch_idx = gl_GlobalInvocationID.y; +#endif + +#ifndef MUL_MAT_ID + uint batch_idx_a = 0; + if (batch_idx != 0) { + const uint i13 = batch_idx / p.ne12; + const uint i12 = batch_idx % p.ne12; + + const uint i03 = i13 / p.broadcast3; + const uint i02 = i12 / p.broadcast2; + + batch_idx_a = i03 * p.ne02 + i02; + } +#else + const uint expert_id = data_ids[expert_idx]; +#endif + + a_offset = +#ifdef MUL_MAT_ID + expert_id * p.batch_stride_a; +#else + batch_idx_a * p.batch_stride_a; +#endif + b_offset = +#ifdef MUL_MAT_ID + (expert_idx % p.ne11) * p.stride_b; +#else + batch_idx * p.batch_stride_b; +#endif + d_offset = +#ifdef MUL_MAT_ID + expert_idx * p.stride_d; +#else + batch_idx * p.batch_stride_d; +#endif +} + +layout (constant_id = 0) const uint BLOCK_SIZE = 32; +layout (constant_id = 1) const uint NUM_ROWS = 1; +layout (constant_id = 2) const uint NUM_COLS = 1; + +shared FLOAT_TYPE tmpsh[NUM_COLS][NUM_ROWS][BLOCK_SIZE]; + +void reduce_result(const in FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const in uint32_t d_offset, const in uint32_t first_row, const in uint32_t num_rows, const in uint32_t tid) { + // sum up partial sums and write back result + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + tmpsh[j][n][tid] = temp[j][n]; + } + } + barrier(); + [[unroll]] for (uint s = BLOCK_SIZE/2; s > 0; s >>= 1) { + if (tid < s) { + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + tmpsh[j][n][tid] += tmpsh[j][n][tid + s]; + } + } + } + barrier(); + } + if (tid == 0) { + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + data_d[j*p.batch_stride_d + d_offset + first_row + n] = D_TYPE(tmpsh[j][n][0]); + } + } + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp new file mode 100644 index 000000000..1cc4996d3 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp @@ -0,0 +1,71 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require + +#define BLOCK_SIZE 32 +#define FLOAT_TYPE float + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE dst[];}; + +layout (push_constant) uniform parameter +{ + uint ncols_x; + uint nrows_x; + uint row_stride_x; + uint channel_stride_x; + uint channel_x_divisor; + uint b_offset; + uint d_offset; +} p; + +shared FLOAT_TYPE tmp[BLOCK_SIZE]; + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint row_x = gl_GlobalInvocationID.y; + const uint channel = gl_GlobalInvocationID.z; + const uint channel_x = channel / p.channel_x_divisor; + + const uint nrows_y = p.ncols_x; + const uint nrows_dst = p.nrows_x; + const uint row_dst = row_x; + + const uint idst = channel*nrows_dst + row_dst; + + tmp[tid] = 0.0f; + + for (uint col_x0 = 0; col_x0 < p.ncols_x; col_x0 += BLOCK_SIZE) { + const uint col_x = col_x0 + tid; + + if (col_x >= p.ncols_x) { + break; + } + + const uint row_y = col_x; + + const uint ix = channel_x*p.channel_stride_x + row_x*p.row_stride_x + col_x; + const uint iy = channel*nrows_y + row_y; + + const FLOAT_TYPE xi = FLOAT_TYPE(data_a[ix]); + + tmp[tid] = fma(xi, FLOAT_TYPE(data_b[iy]), tmp[tid]); + } + + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] += tmp[tid + s]; + } + barrier(); + } + + if (tid == 0) { + dst[idst] = tmp[0]; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp new file mode 100644 index 000000000..9b443807d --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp @@ -0,0 +1,73 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require + +#define BLOCK_SIZE 32 +#define FLOAT_TYPE float + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE dst[];}; + +layout (push_constant) uniform parameter +{ + uint ncols_x; + uint nrows_x; + uint nchannels_x; + uint nchannels_y; + uint b_offset; + uint d_offset; +} p; + +shared FLOAT_TYPE tmp[BLOCK_SIZE]; + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint row_x = gl_GlobalInvocationID.y; + const uint channel = gl_GlobalInvocationID.z; + const uint channel_x = channel / (p.nchannels_y / p.nchannels_x); + + const uint nrows_y = p.ncols_x; + const uint nrows_dst = p.nrows_x; + const uint row_dst = row_x; + + tmp[tid] = FLOAT_TYPE(0.0f); + + for (uint col_x0 = 0; col_x0 < p.ncols_x; col_x0 += BLOCK_SIZE) { + const uint col_x = col_x0 + tid; + + if (col_x >= p.ncols_x) { + break; + } + + // x is transposed and permuted + const uint ix = row_x*p.nchannels_x*p.ncols_x + channel_x*p.ncols_x + col_x; + const FLOAT_TYPE xi = FLOAT_TYPE(data_a[ix]); + + const uint row_y = col_x; + + // y is not transposed but permuted + const uint iy = channel*nrows_y + row_y; + + tmp[tid] = fma(xi, FLOAT_TYPE(data_b[iy]), tmp[tid]); + } + + // dst is not transposed and not permuted + const uint idst = channel*nrows_dst + row_dst; + + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] += tmp[tid + s]; + } + barrier(); + } + + if (tid == 0) { + dst[idst] = tmp[0]; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp new file mode 100644 index 000000000..934213446 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp @@ -0,0 +1,115 @@ +#version 450 +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads are used to process each block + const uint it_size = gl_WorkGroupSize.x/16; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid%16; // 0...16 + const uint ix = tid/16; + + const uint step = 8; + + const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... + const uint v_in = itid - step*v_im; // 0...15 or 0...7 + + const uint l0 = 2*v_in; // 0...15 + const uint q_offset = 32*v_im + l0; + const uint s_offset = 8*v_im; + const uint y_offset = 128*v_im + l0; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { + const uint y_idx = i * QUANT_K + y_offset; + + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; + f16vec2 d = data_a[ib0 + i].d; + const FLOAT_TYPE dall = d.x; + const FLOAT_TYPE dmin = d.y; + + uint32_t s0_u32 = data_a_packed32[ib0 + i].scales[s_offset / 4 + 0]; + uint32_t s4_u32 = data_a_packed32[ib0 + i].scales[s_offset / 4 + 1]; + + uint32_t s0_lo4_u32 = s0_u32 & 0x0F0F0F0F; + uint32_t s0_hi4_u32 = (s0_u32 >> 4) & 0x0F0F0F0F; + uint32_t s4_lo4_u32 = s4_u32 & 0x0F0F0F0F; + uint32_t s4_hi4_u32 = (s4_u32 >> 4) & 0x0F0F0F0F; + + uvec4 s0_lo4 = uvec4(unpack8(s0_lo4_u32)); + uvec4 s4_lo4 = uvec4(unpack8(s4_lo4_u32)); + uvec4 s0_hi4 = uvec4(unpack8(s0_hi4_u32)); + uvec4 s4_hi4 = uvec4(unpack8(s4_hi4_u32)); + + uint16_t qs0_u16 = data_a_packed16[ib0 + i].qs[q_offset / 2 + 0]; + uint16_t qs16_u16 = data_a_packed16[ib0 + i].qs[q_offset / 2 + 8]; + uvec2 qs0 = uvec2(unpack8(qs0_u16)); + uvec2 qs16 = uvec2(unpack8(qs16_u16)); + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + B_TYPE_VEC2 b0 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 0]; + B_TYPE_VEC2 b16 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 8]; + B_TYPE_VEC2 b32 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 16]; + B_TYPE_VEC2 b48 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 24]; + B_TYPE_VEC2 b64 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 32]; + B_TYPE_VEC2 b80 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 40]; + B_TYPE_VEC2 b96 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 48]; + B_TYPE_VEC2 b112 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 56]; + + FLOAT_TYPE sum1 = FLOAT_TYPE(0.0); + FLOAT_TYPE sum2 = FLOAT_TYPE(0.0); + [[unroll]] for (int l = 0; l < 2; ++l) { + sum1 = fma(FLOAT_TYPE(b0[l]), FLOAT_TYPE(s0_lo4[0]) * FLOAT_TYPE((qs0[l] >> 0) & 3), + fma(FLOAT_TYPE(b16[l]), FLOAT_TYPE(s0_lo4[1]) * FLOAT_TYPE((qs16[l] >> 0) & 3), + fma(FLOAT_TYPE(b32[l]), FLOAT_TYPE(s0_lo4[2]) * FLOAT_TYPE((qs0[l] >> 2) & 3), + fma(FLOAT_TYPE(b48[l]), FLOAT_TYPE(s0_lo4[3]) * FLOAT_TYPE((qs16[l] >> 2) & 3), + fma(FLOAT_TYPE(b64[l]), FLOAT_TYPE(s4_lo4[0]) * FLOAT_TYPE((qs0[l] >> 4) & 3), + fma(FLOAT_TYPE(b80[l]), FLOAT_TYPE(s4_lo4[1]) * FLOAT_TYPE((qs16[l] >> 4) & 3), + fma(FLOAT_TYPE(b96[l]), FLOAT_TYPE(s4_lo4[2]) * FLOAT_TYPE((qs0[l] >> 6) & 3), + fma(FLOAT_TYPE(b112[l]), FLOAT_TYPE(s4_lo4[3]) * FLOAT_TYPE((qs16[l] >> 6) & 3), sum1)))))))); + sum2 = fma(FLOAT_TYPE(b0[l]), FLOAT_TYPE(s0_hi4[0]), + fma(FLOAT_TYPE(b16[l]), FLOAT_TYPE(s0_hi4[1]), + fma(FLOAT_TYPE(b32[l]), FLOAT_TYPE(s0_hi4[2]), + fma(FLOAT_TYPE(b48[l]), FLOAT_TYPE(s0_hi4[3]), + fma(FLOAT_TYPE(b64[l]), FLOAT_TYPE(s4_hi4[0]), + fma(FLOAT_TYPE(b80[l]), FLOAT_TYPE(s4_hi4[1]), + fma(FLOAT_TYPE(b96[l]), FLOAT_TYPE(s4_hi4[2]), + fma(FLOAT_TYPE(b112[l]), FLOAT_TYPE(s4_hi4[3]), sum2)))))))); + } + temp[j][n] = fma(dall, sum1, fma(-dmin, sum2, temp[j][n])); + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp new file mode 100644 index 000000000..86b0159d9 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp @@ -0,0 +1,103 @@ +#version 450 +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads are used to process each block + const uint it_size = gl_WorkGroupSize.x/16; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid%16; // 0...16 + const uint ix = tid/16; + + const uint step = 8; + + const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... + const uint v_in = itid - step*v_im; // 0...15 or 0...7 + + const uint8_t m = uint8_t(1 << (4 * v_im)); + + const uint l0 = 2*v_in; // 0...15 + const uint q_offset = 32*v_im + l0; + const uint y_offset = 128*v_im + l0; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + const uint s_shift = 4 * v_im; + + [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { + const uint y_idx = i * QUANT_K + y_offset; + + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; + const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d); + + uint16_t s0_16 = data_a_packed16[ib0 + i].scales[0]; + uint16_t s2_16 = data_a_packed16[ib0 + i].scales[1]; + uint16_t s4_16 = data_a_packed16[ib0 + i].scales[2]; + uint16_t s6_16 = data_a_packed16[ib0 + i].scales[3]; + uint16_t s8_16 = data_a_packed16[ib0 + i].scales[4]; + uint16_t s10_16 = data_a_packed16[ib0 + i].scales[5]; + u8vec2 s0 = unpack8(s0_16); + u8vec2 s2 = unpack8(s2_16); + u8vec2 s4 = unpack8(s4_16); + u8vec2 s6 = unpack8(s6_16); + u8vec2 s8 = unpack8(s8_16); + u8vec2 s10 = unpack8(s10_16); + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + + B_TYPE_VEC2 b0 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 0]; + B_TYPE_VEC2 b16 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 8]; + B_TYPE_VEC2 b32 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 16]; + B_TYPE_VEC2 b48 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 24]; + B_TYPE_VEC2 b64 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 32]; + B_TYPE_VEC2 b80 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 40]; + B_TYPE_VEC2 b96 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 48]; + B_TYPE_VEC2 b112 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 56]; + + FLOAT_TYPE sum = FLOAT_TYPE(0.0); + [[unroll]] for (int l = 0; l < 2; ++l) { + sum = fma(FLOAT_TYPE(b0[l]) * FLOAT_TYPE(int8_t(((s0[0] >> s_shift) & 0xF) | ((s8[0] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 0)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b32[l]) * FLOAT_TYPE(int8_t(((s2[0] >> s_shift) & 0xF) | ((s10[0] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 1)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b64[l]) * FLOAT_TYPE(int8_t(((s4[0] >> s_shift) & 0xF) | ((s8[0] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 2)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b96[l]) * FLOAT_TYPE(int8_t(((s6[0] >> s_shift) & 0xF) | ((s10[0] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 3)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b16[l]) * FLOAT_TYPE(int8_t(((s0[1] >> s_shift) & 0xF) | ((s8[1] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 0)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b48[l]) * FLOAT_TYPE(int8_t(((s2[1] >> s_shift) & 0xF) | ((s10[1] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 1)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b80[l]) * FLOAT_TYPE(int8_t(((s4[1] >> s_shift) & 0xF) | ((s8[1] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 2)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b112[l]) * FLOAT_TYPE(int8_t(((s6[1] >> s_shift) & 0xF) | ((s10[1] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 3)) != 0) ? 0 : 4)), sum)))))))); + } + temp[j][n] = fma(d, sum, temp[j][n]); + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp new file mode 100644 index 000000000..cd1dd8e89 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp @@ -0,0 +1,133 @@ +#version 450 + +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads are used to process each block + const uint it_size = gl_WorkGroupSize.x/16; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid%16; // 0...16 + const uint ix = tid/16; + + const uint step = 4; + + const uint il = itid/step; // 0...3 + const uint ir = itid - step*il; // 0...7 or 0...3 + const uint n = 4; + + const uint v_im = il / 2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224 + const uint v_in = il % 2; + + const uint l0 = n * (2 * ir + v_in); // 0...15 + const uint q_offset = 32*v_im + l0; + const uint y_offset = 64*v_im + l0; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { + const uint y1_idx = i * QUANT_K + y_offset; + const uint y2_idx = y1_idx + 128; + + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; + f16vec2 d = data_a[ib0 + i].d; + const FLOAT_TYPE dall = FLOAT_TYPE(d.x); + const FLOAT_TYPE dmin = FLOAT_TYPE(d.y); + + uint32_t scale0_u32 = data_a_packed16[ib0 + i].scales[v_im ]; + uint32_t scale4_u32 = data_a_packed16[ib0 + i].scales[v_im + 2]; + uint32_t scale8_u32 = data_a_packed16[ib0 + i].scales[v_im + 4]; + uvec4 scale0 = uvec4(unpack8(scale0_u32)); + uvec4 scale4 = uvec4(unpack8(scale4_u32)); + uvec4 scale8 = uvec4(unpack8(scale8_u32)); + + const uint32_t sc0 = ( scale0.x & 0x3f); + const uint32_t sc1 = ( scale0.y & 0x3f); + const uint32_t sc2 = ( scale4.x & 0x3f); + const uint32_t sc3 = ( scale4.y & 0x3f); + const uint32_t sc4 = (( scale8.x & 0x0f) | ((scale0.x & 0xc0) >> 2)); + const uint32_t sc5 = (( scale8.y & 0x0f) | ((scale0.y & 0xc0) >> 2)); + const uint32_t sc6 = (((scale8.x >> 4) & 0x0f) | ((scale4.x & 0xc0) >> 2)); + const uint32_t sc7 = (((scale8.y >> 4) & 0x0f) | ((scale4.y & 0xc0) >> 2)); + + uint32_t qs0_u32 = data_a_packed32[ib0 + i].qs[q_offset / 4]; + uint32_t qs64_u32 = data_a_packed32[ib0 + i].qs[q_offset / 4 + 16]; + + uint32_t qs0_u32_lo4 = qs0_u32 & 0x0F0F0F0F; + uint32_t qs0_u32_hi4 = (qs0_u32 >> 4) & 0x0F0F0F0F; + uint32_t qs64_u32_lo4 = qs64_u32 & 0x0F0F0F0F; + uint32_t qs64_u32_hi4 = (qs64_u32 >> 4) & 0x0F0F0F0F; + + uvec4 qs0_lo4 = uvec4(unpack8(qs0_u32_lo4)); + uvec4 qs64_lo4 = uvec4(unpack8(qs64_u32_lo4)); + uvec4 qs0_hi4 = uvec4(unpack8(qs0_u32_hi4)); + uvec4 qs64_hi4 = uvec4(unpack8(qs64_u32_hi4)); + + const uint32_t q4_0 = qs0_lo4.x; + const uint32_t q4_1 = qs0_lo4.y; + const uint32_t q4_2 = qs0_lo4.z; + const uint32_t q4_3 = qs0_lo4.w; + const uint32_t q4_4 = qs0_hi4.x; + const uint32_t q4_5 = qs0_hi4.y; + const uint32_t q4_6 = qs0_hi4.z; + const uint32_t q4_7 = qs0_hi4.w; + const uint32_t q4_8 = qs64_lo4.x; + const uint32_t q4_9 = qs64_lo4.y; + const uint32_t q4_10 = qs64_lo4.z; + const uint32_t q4_11 = qs64_lo4.w; + const uint32_t q4_12 = qs64_hi4.x; + const uint32_t q4_13 = qs64_hi4.y; + const uint32_t q4_14 = qs64_hi4.z; + const uint32_t q4_15 = qs64_hi4.w; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + B_TYPE_VEC4 by10 = data_b_v4[(j*p.batch_stride_b + b_offset + y1_idx) / 4]; + B_TYPE_VEC4 by132 = data_b_v4[(j*p.batch_stride_b + b_offset + y1_idx) / 4 + 8]; + B_TYPE_VEC4 by20 = data_b_v4[(j*p.batch_stride_b + b_offset + y2_idx) / 4]; + B_TYPE_VEC4 by232 = data_b_v4[(j*p.batch_stride_b + b_offset + y2_idx) / 4 + 8]; + + const FLOAT_TYPE sx = fma(FLOAT_TYPE(by10.x), q4_0, fma(FLOAT_TYPE(by10.y), q4_1, fma(FLOAT_TYPE(by10.z), q4_2, FLOAT_TYPE(by10.w) * q4_3))); + const FLOAT_TYPE sy = fma(FLOAT_TYPE(by132.x), q4_4, fma(FLOAT_TYPE(by132.y), q4_5, fma(FLOAT_TYPE(by132.z), q4_6, FLOAT_TYPE(by132.w) * q4_7))); + const FLOAT_TYPE sz = fma(FLOAT_TYPE(by20.x), q4_8, fma(FLOAT_TYPE(by20.y), q4_9, fma(FLOAT_TYPE(by20.z), q4_10, FLOAT_TYPE(by20.w) * q4_11))); + const FLOAT_TYPE sw = fma(FLOAT_TYPE(by232.x), q4_12, fma(FLOAT_TYPE(by232.y), q4_13, fma(FLOAT_TYPE(by232.z), q4_14, FLOAT_TYPE(by232.w) * q4_15))); + const FLOAT_TYPE smin = + fma(FLOAT_TYPE(by10.x), sc2, fma(FLOAT_TYPE(by132.x), sc3, fma(FLOAT_TYPE(by20.x), sc6, fma(FLOAT_TYPE(by232.x), sc7, + fma(FLOAT_TYPE(by10.y), sc2, fma(FLOAT_TYPE(by132.y), sc3, fma(FLOAT_TYPE(by20.y), sc6, fma(FLOAT_TYPE(by232.y), sc7, + fma(FLOAT_TYPE(by10.z), sc2, fma(FLOAT_TYPE(by132.z), sc3, fma(FLOAT_TYPE(by20.z), sc6, fma(FLOAT_TYPE(by232.z), sc7, + fma(FLOAT_TYPE(by10.w), sc2, fma(FLOAT_TYPE(by132.w), sc3, fma(FLOAT_TYPE(by20.w), sc6, FLOAT_TYPE(by232.w) * sc7))))))))))))))); + temp[j][n] = fma(dall, fma(sx, sc0, fma(sy, sc1, fma(sz, sc4, sw * sc5))), fma(-dmin, smin, temp[j][n])); + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp new file mode 100644 index 000000000..0a68891c3 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp @@ -0,0 +1,162 @@ +#version 450 + +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads are used to process each block + const uint it_size = gl_WorkGroupSize.x/16; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid%16; // 0...16 + const uint ix = tid/16; + + const uint il = itid/4; // 0...3 + const uint ir = itid - 4*il; // 0...7 or 0...3 + + const uint v_im = il / 2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224 + const uint v_in = il % 2; + + const uint l0 = 4*ir + 2*v_in; // 0...15 + const uint q_offset = 32*v_im + l0; + const uint y_offset = 64*v_im + l0; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { + const uint y1_idx = i * QUANT_K + y_offset; + const uint y2_idx = y1_idx + 128; + + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; + f16vec2 d = data_a[ib0 + i].d; + const FLOAT_TYPE dall = FLOAT_TYPE(d.x); + const FLOAT_TYPE dmin = FLOAT_TYPE(d.y); + + uint32_t scale0_u32 = data_a_packed16[ib0 + i].scales[v_im ]; + uint32_t scale4_u32 = data_a_packed16[ib0 + i].scales[v_im + 2]; + uint32_t scale8_u32 = data_a_packed16[ib0 + i].scales[v_im + 4]; + uvec4 scale0 = uvec4(unpack8(scale0_u32)); + uvec4 scale4 = uvec4(unpack8(scale4_u32)); + uvec4 scale8 = uvec4(unpack8(scale8_u32)); + + const uint32_t sc0 = ( scale0.x & 0x3f); + const uint32_t sc1 = ( scale0.y & 0x3f); + const uint32_t sc2 = ( scale4.x & 0x3f); + const uint32_t sc3 = ( scale4.y & 0x3f); + const uint32_t sc4 = (( scale8.x & 0x0f) | ((scale0.x & 0xc0) >> 2)); + const uint32_t sc5 = (( scale8.y & 0x0f) | ((scale0.y & 0xc0) >> 2)); + const uint32_t sc6 = (((scale8.x >> 4) & 0x0f) | ((scale4.x & 0xc0) >> 2)); + const uint32_t sc7 = (((scale8.y >> 4) & 0x0f) | ((scale4.y & 0xc0) >> 2)); + + uint32_t qs0_16_u32 = uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 8]) << 16); + uint32_t qs64_80_u32 = uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 32]) | (uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 40]) << 16); + + uint32_t qs0_16_u32_lo4 = qs0_16_u32 & 0x0F0F0F0F; + uint32_t qs0_16_u32_hi4 = (qs0_16_u32 >> 4) & 0x0F0F0F0F; + uint32_t qs64_80_u32_lo4 = qs64_80_u32 & 0x0F0F0F0F; + uint32_t qs64_80_u32_hi4 = (qs64_80_u32 >> 4) & 0x0F0F0F0F; + + uint32_t qh = pack32(u16vec2(data_a_packed16[ib0 + i].qh[l0 / 2], data_a_packed16[ib0 + i].qh[l0 / 2 + 8])); + + uint32_t qs0_16_lo4_offset16 = ((qh >> (2*v_im)) & 0x01010101) << 4; + uint32_t qs0_16_hi4_offset16 = ((qh >> (2*v_im)) & 0x02020202) << 3; + uint32_t qs64_80_lo4_offset16 = ((qh >> (2*v_im)) & 0x10101010) << 0; + uint32_t qs64_80_hi4_offset16 = ((qh >> (2*v_im)) & 0x20202020) >> 1; + + qs0_16_u32_lo4 += qs0_16_lo4_offset16; + qs0_16_u32_hi4 += qs0_16_hi4_offset16; + qs64_80_u32_lo4 += qs64_80_lo4_offset16; + qs64_80_u32_hi4 += qs64_80_hi4_offset16; + + uvec4 qs0_16_lo4 = uvec4(unpack8(qs0_16_u32_lo4)); + uvec4 qs64_80_lo4 = uvec4(unpack8(qs64_80_u32_lo4)); + uvec4 qs0_16_hi4 = uvec4(unpack8(qs0_16_u32_hi4)); + uvec4 qs64_80_hi4 = uvec4(unpack8(qs64_80_u32_hi4)); + + const uint32_t q4_0 = qs0_16_lo4.x; + const uint32_t q4_1 = qs0_16_lo4.y; + const uint32_t q4_2 = qs0_16_lo4.z; + const uint32_t q4_3 = qs0_16_lo4.w; + const uint32_t q4_4 = qs0_16_hi4.x; + const uint32_t q4_5 = qs0_16_hi4.y; + const uint32_t q4_6 = qs0_16_hi4.z; + const uint32_t q4_7 = qs0_16_hi4.w; + const uint32_t q4_8 = qs64_80_lo4.x; + const uint32_t q4_9 = qs64_80_lo4.y; + const uint32_t q4_10 = qs64_80_lo4.z; + const uint32_t q4_11 = qs64_80_lo4.w; + const uint32_t q4_12 = qs64_80_hi4.x; + const uint32_t q4_13 = qs64_80_hi4.y; + const uint32_t q4_14 = qs64_80_hi4.z; + const uint32_t q4_15 = qs64_80_hi4.w; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + B_TYPE_VEC2 by10 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2]; + B_TYPE_VEC2 by116 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 8]; + B_TYPE_VEC2 by132 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 16]; + B_TYPE_VEC2 by148 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 24]; + B_TYPE_VEC2 by20 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2]; + B_TYPE_VEC2 by216 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 8]; + B_TYPE_VEC2 by232 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 16]; + B_TYPE_VEC2 by248 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 24]; + + const FLOAT_TYPE sx = + fma(FLOAT_TYPE(by10.x), q4_0, + fma(FLOAT_TYPE(by10.y), q4_1, + fma(FLOAT_TYPE(by116.x), q4_2, + FLOAT_TYPE(by116.y) * q4_3))); + const FLOAT_TYPE sy = + fma(FLOAT_TYPE(by132.x), q4_4, + fma(FLOAT_TYPE(by132.y), q4_5, + fma(FLOAT_TYPE(by148.x), q4_6, + FLOAT_TYPE(by148.y) * q4_7))); + const FLOAT_TYPE sz = + fma(FLOAT_TYPE(by20.x), q4_8, + fma(FLOAT_TYPE(by20.y), q4_9, + fma(FLOAT_TYPE(by216.x), q4_10, + FLOAT_TYPE(by216.y) * q4_11))); + const FLOAT_TYPE sw = + fma(FLOAT_TYPE(by232.x), q4_12, + fma(FLOAT_TYPE(by232.y), q4_13, + fma(FLOAT_TYPE(by248.x), q4_14, + FLOAT_TYPE(by248.y) * q4_15))); + const FLOAT_TYPE smin = + fma(FLOAT_TYPE(by10.x) + FLOAT_TYPE(by10.y) + FLOAT_TYPE(by116.x) + FLOAT_TYPE(by116.y), sc2, + fma(FLOAT_TYPE(by132.x) + FLOAT_TYPE(by132.y) + FLOAT_TYPE(by148.x) + FLOAT_TYPE(by148.y), sc3, + fma(FLOAT_TYPE(by20.x) + FLOAT_TYPE(by20.y) + FLOAT_TYPE(by216.x) + FLOAT_TYPE(by216.y), sc6, + (FLOAT_TYPE(by232.x) + FLOAT_TYPE(by232.y) + FLOAT_TYPE(by248.x) + FLOAT_TYPE(by248.y)) * sc7))); + temp[j][n] = fma(dall, fma(sx, sc0, fma(sy, sc1, fma(sz, sc4, sw * sc5))), fma(-dmin, smin, temp[j][n])); + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp new file mode 100644 index 000000000..70e13a56b --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp @@ -0,0 +1,112 @@ +#version 450 + +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads are used to process each block + const uint it_size = gl_WorkGroupSize.x/16; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid%16; // 0...16 + const uint ix = tid/16; + + const uint step = 8; + + const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... + const uint v_in = itid - step*v_im; // 0...15 or 0...7 + + const uint l0 = 4 * v_in; // 0, 4, 8, ..., 28 + const uint is = v_in / 4; + + const uint ql_offset = 64*v_im + l0; + const uint qh_offset = 32*v_im + l0; + const uint s_offset = 8*v_im + is; + const uint y_offset = 128*v_im + l0; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { + const uint y_idx = i * QUANT_K + y_offset; + + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; + const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d); + + FLOAT_TYPE scales[4]; + scales[0] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 0]); + scales[1] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 2]); + scales[2] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 4]); + scales[3] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 6]); + + uint32_t ql0_u32 = uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 1]) << 16); + uint32_t ql32_u32 = uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 16]) | (uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 17]) << 16); + + uint32_t ql0_u32_lo4 = ql0_u32 & 0x0F0F0F0F; + uint32_t ql0_u32_hi4 = (ql0_u32 >> 4) & 0x0F0F0F0F; + uint32_t ql32_u32_lo4 = ql32_u32 & 0x0F0F0F0F; + uint32_t ql32_u32_hi4 = (ql32_u32 >> 4) & 0x0F0F0F0F; + + uint32_t qh_u32 = uint32_t(data_a_packed16[ib0 + i].qh[qh_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].qh[qh_offset / 2 + 1]) << 16); + uint32_t qh0_u32 = (qh_u32 & 0x03030303) << 4; + uint32_t qh2_u32 = (qh_u32 & 0x0C0C0C0C) << 2; + uint32_t qh4_u32 = (qh_u32 & 0x30303030) << 0; + uint32_t qh6_u32 = (qh_u32 & 0xC0C0C0C0) >> 2; + + uint32_t q0_u32 = ql0_u32_lo4 | qh0_u32; + uint32_t q1_u32 = ql32_u32_lo4 | qh2_u32; + uint32_t q2_u32 = ql0_u32_hi4 | qh4_u32; + uint32_t q3_u32 = ql32_u32_hi4 | qh6_u32; + + uvec4 q0 = uvec4(unpack8(q0_u32)); + uvec4 q1 = uvec4(unpack8(q1_u32)); + uvec4 q2 = uvec4(unpack8(q2_u32)); + uvec4 q3 = uvec4(unpack8(q3_u32)); + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + B_TYPE_VEC4 by0 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4]; + B_TYPE_VEC4 by32 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 8]; + B_TYPE_VEC4 by64 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 16]; + B_TYPE_VEC4 by96 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 24]; + + FLOAT_TYPE sum = FLOAT_TYPE(0.0); + [[unroll]] for (int l = 0; l < 4; ++l) { + sum = fma(FLOAT_TYPE(by0[l]) * scales[0], FLOAT_TYPE(int8_t(q0[l]) - 32), + fma(FLOAT_TYPE(by32[l]) * scales[1], FLOAT_TYPE(int8_t(q1[l]) - 32), + fma(FLOAT_TYPE(by64[l]) * scales[2], FLOAT_TYPE(int8_t(q2[l]) - 32), + fma(FLOAT_TYPE(by96[l]) * scales[3], FLOAT_TYPE(int8_t(q3[l]) - 32), sum)))); + } + temp[j][n] += sum * d; + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp new file mode 100644 index 000000000..48122cbef --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -0,0 +1,631 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require + +#ifdef FLOAT16 +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#endif + +#ifdef COOPMAT +#extension GL_KHR_cooperative_matrix : enable +#extension GL_KHR_memory_scope_semantics : enable +#extension GL_KHR_shader_subgroup_basic : enable +#endif + +#ifdef MUL_MAT_ID +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#endif + +#include "types.comp" + +#ifndef LOAD_VEC_A +#define LOAD_VEC_A 1 +#endif +#ifndef LOAD_VEC_B +#define LOAD_VEC_B 1 +#endif + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; + +#ifdef MUL_MAT_ID +layout (binding = 3) readonly buffer IDS {int data_ids[];}; +#endif + +layout (push_constant) uniform parameter +{ + uint M; + uint N; + uint K; + uint stride_a; + uint stride_b; + uint stride_d; + + uint batch_stride_a; + uint batch_stride_b; + uint batch_stride_d; + +#ifdef MUL_MAT_ID + uint nei0; + uint nei1; + uint nbi1; + uint ne11; +#else + uint k_split; + uint ne02; + uint ne12; + uint broadcast2; + uint broadcast3; +#endif +} p; + +layout (constant_id = 0) const uint BLOCK_SIZE = 64; +layout (constant_id = 1) const uint BM = 64; +layout (constant_id = 2) const uint BN = 64; +layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working with a quant +layout (constant_id = 4) const uint WM = 32; +layout (constant_id = 5) const uint WN = 32; +layout (constant_id = 6) const uint WMITER = 2; +layout (constant_id = 7) const uint TM = 4; +layout (constant_id = 8) const uint TN = 2; +layout (constant_id = 9) const uint TK = 1; // Only needed for coopmat +layout (constant_id = 10) const uint WARP = 32; + +#ifdef COOPMAT +#define SHMEM_STRIDE (BK + 8) +#else +#define SHMEM_STRIDE (BK + 1) +#endif + +shared FLOAT_TYPE buf_a[BM * SHMEM_STRIDE]; +shared FLOAT_TYPE buf_b[BN * SHMEM_STRIDE]; + +#ifdef MUL_MAT_ID +shared u16vec2 row_ids[3072]; +#endif // MUL_MAT_ID + +#define NUM_WARPS (BLOCK_SIZE / WARP) + +#ifdef COOPMAT +shared ACC_TYPE coopmat_stage[TM * TN * NUM_WARPS]; +#endif + +void main() { +#if defined(DATA_A_IQ4_NL) + init_iq4nl_shmem(); +#endif + +#ifdef MUL_MAT_ID + const uint expert_idx = gl_GlobalInvocationID.z; +#else + const uint batch_idx = gl_GlobalInvocationID.z; + + const uint i13 = batch_idx / p.ne12; + const uint i12 = batch_idx % p.ne12; + + const uint i03 = i13 / p.broadcast3; + const uint i02 = i12 / p.broadcast2; + + const uint batch_idx_a = i03 * p.ne02 + i02; +#endif + + const uint blocks_m = (p.M + BM - 1) / BM; + const uint ir = gl_WorkGroupID.x % blocks_m; + const uint ik = gl_WorkGroupID.x / blocks_m; + const uint ic = gl_WorkGroupID.y; + + const uint WNITER = (WM * WN) / (WARP * TM * TN * WMITER); + const uint WSUBM = WM / WMITER; + const uint WSUBN = WN / WNITER; + +#ifdef COOPMAT + const uint warp_i = gl_SubgroupID; + + const uint tiw = gl_SubgroupInvocationID; + + const uint cms_per_row = WM / TM; + const uint cms_per_col = WN / TN; + + const uint storestride = WARP / TM; + const uint store_r = tiw % TM; + const uint store_c = tiw / TM; +#else + const uint warp_i = gl_LocalInvocationID.x / WARP; + + const uint tiw = gl_LocalInvocationID.x % WARP; + + const uint tiwr = tiw % (WSUBM / TM); + const uint tiwc = tiw / (WSUBM / TM); +#endif + + const uint warp_r = warp_i % (BM / WM); + const uint warp_c = warp_i / (BM / WM); + + const uint loadr_a = gl_LocalInvocationID.x % (BK / LOAD_VEC_A); + const uint loadc_a = gl_LocalInvocationID.x / (BK / LOAD_VEC_A); + const uint loadr_b = gl_LocalInvocationID.x % (BK / LOAD_VEC_B); + const uint loadc_b = gl_LocalInvocationID.x / (BK / LOAD_VEC_B); + + const uint loadstride_a = gl_WorkGroupSize.x * LOAD_VEC_A / BK; + const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B / BK; + +#ifdef MUL_MAT_ID + uint _ne1 = 0; + for (uint ii1 = 0; ii1 < p.nei1; ii1++) { + for (uint ii0 = 0; ii0 < p.nei0; ii0++) { + if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { + row_ids[_ne1] = u16vec2(ii0, ii1); + _ne1++; + } + } + } + + barrier(); + + // Workgroup has no work + if (ic * BN >= _ne1) return; +#endif + +#ifdef MUL_MAT_ID + const uint start_k = 0; + const uint end_k = p.K; +#else + const uint start_k = ik * p.k_split; + const uint end_k = min(p.K, (ik + 1) * p.k_split); +#endif + + uint pos_a = ( +#ifdef MUL_MAT_ID + expert_idx * p.batch_stride_a + +#else + batch_idx_a * p.batch_stride_a + +#endif + ir * BM * p.stride_a + start_k) / LOAD_VEC_A; +#ifdef MUL_MAT_ID + uint pos_b = 0; +#else + uint pos_b = (batch_idx * p.batch_stride_b + ic * BN * p.stride_b + start_k) / LOAD_VEC_B; +#endif + +#ifdef COOPMAT + coopmat cache_a; + coopmat cache_b; + coopmat sums[cms_per_row * cms_per_col]; + + [[unroll]] for (uint i = 0; i < cms_per_row * cms_per_col; i++) { + sums[i] = coopmat(0.0f); + } +#else + ACC_TYPE sums[WMITER * TM * WNITER * TN]; + FLOAT_TYPE cache_a[WMITER * TM]; + FLOAT_TYPE cache_b[WNITER * TN]; + + [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN; i++) { + sums[i] = ACC_TYPE(0.0f); + } +#endif + + for (uint block = start_k; block < end_k; block += BK) { + [[unroll]] for (uint l = 0; l < BM; l += loadstride_a) { + +#if defined(DATA_A_F32) || defined(DATA_A_F16) +#if LOAD_VEC_A == 8 + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx][0].x); + buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx][0].y); + buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx][0].z); + buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx][0].w); + buf_a[buf_idx + 4] = FLOAT_TYPE(data_a[idx][1].x); + buf_a[buf_idx + 5] = FLOAT_TYPE(data_a[idx][1].y); + buf_a[buf_idx + 6] = FLOAT_TYPE(data_a[idx][1].z); + buf_a[buf_idx + 7] = FLOAT_TYPE(data_a[idx][1].w); +#elif LOAD_VEC_A == 4 + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx].x); + buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx].y); + buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx].z); + buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx].w); +#else + if (ir * BM + loadc_a + l < p.M && block + loadr_a < end_k) { + buf_a[(loadc_a + l) * SHMEM_STRIDE + loadr_a] = FLOAT_TYPE(data_a[pos_a + (loadc_a + l) * p.stride_a + loadr_a]); + } else { + buf_a[(loadc_a + l) * SHMEM_STRIDE + loadr_a] = FLOAT_TYPE(0.0f); + } +#endif +#elif defined(DATA_A_Q4_0) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; + + const uint ib = idx / 16; + const uint iqs = idx & 0xF; + + const float d = float(data_a[ib].d); + const uint vui = uint(data_a[ib].qs[iqs]); + const vec2 v = (vec2(vui & 0xF, vui >> 4) - 8.0f) * d; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q4_1) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; + + const uint ib = idx / 16; + const uint iqs = idx & 0xF; + + const float d = float(data_a[ib].d); + const float m = float(data_a[ib].m); + const uint vui = uint(data_a[ib].qs[iqs]); + const vec2 v = vec2(vui & 0xF, vui >> 4) * d + m; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q5_0) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; + + const uint ib = idx / 16; + const uint iqs = idx & 0xF; + + const float d = float(data_a[ib].d); + const uint uint_qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0]; + const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const uint vui = uint(data_a[ib].qs[iqs]); + const vec2 v = (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f) * d; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q5_1) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; + + const uint ib = idx / 16; + const uint iqs = idx & 0xF; + + const float d = float(data_a[ib].d); + const float m = float(data_a[ib].m); + const uint uint_qh = data_a[ib].qh; + const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const uint vui = uint(data_a[ib].qs[iqs]); + const vec2 v = vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) * d + m; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q8_0) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 16; + const uint iqs = (idx & 0xF) * 2; + + const float d = float(data_a[ib].d); + const vec2 v = vec2(int(data_a[ib].qs[iqs]), int(data_a[ib].qs[iqs + 1])) * d; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 1] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q2_K) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint qsi = (iqs / 64) * 32 + (iqs % 16) * 2; // 0,2,4..30 + const uint scalesi = iqs / 8; // 0..15 + const uint qsshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 + + const uvec2 qs = uvec2(data_a[ib].qs[qsi], data_a[ib].qs[qsi + 1]); + const uint scales = data_a[ib].scales[scalesi]; + const vec2 d = vec2(data_a[ib].d); + + const vec2 v = d.x * float(scales & 0xF) * vec2((qs >> qsshift) & 3) - d.y * float(scales >> 4); + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 1] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q3_K) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint n = iqs / 64; // 0,1 + const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 + const uint hmi = (iqs % 16) * 2; // 0,2,4..30 + const uint j = (iqs % 64) / 4; // 0..3 + const uint is = iqs / 8; // 0..15 + const uint halfsplit = ((iqs % 64) / 16); // 0,1,2,3 + const uint qsshift = halfsplit * 2; // 0,2,4,6 + const uint m = 1 << (4 * n + halfsplit); // 1,2,4,8,16,32,64,128 + + const int8_t us = int8_t(is < 4 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+8] >> 0) & 3) << 4) : + is < 8 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+4] >> 2) & 3) << 4) : + is < 12 ? (data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is+0] >> 4) & 3) << 4) : + (data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is-4] >> 6) & 3) << 4)); + const float dl = float(data_a[ib].d) * float(us - 32); + + buf_a[buf_idx ] = FLOAT_TYPE(dl * float(int8_t((data_a[ib].qs[qsi ] >> qsshift) & 3) - (((data_a[ib].hmask[hmi ] & m) != 0) ? 0 : 4))); + buf_a[buf_idx + 1] = FLOAT_TYPE(dl * float(int8_t((data_a[ib].qs[qsi + 1] >> qsshift) & 3) - (((data_a[ib].hmask[hmi + 1] & m) != 0) ? 0 : 4))); +#elif defined(DATA_A_Q4_K) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint n = iqs / 32; // 0,1,2,3 + const uint b = (iqs % 32) / 16; // 0,1 + const uint is = 2 * n + b; // 0..7 + const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 + + const vec2 loadd = vec2(data_a[ib].d); + + const uint scidx0 = (is < 4) ? is : (is + 4); + const uint scidx1 = (is < 4) ? is : (is - 4); + const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; + const uint scidxshift1 = (is < 4) ? 0 : 2; + const uint mbidx0 = is + 4; + const uint mbidx1 = (is < 4) ? is + 4 : is; + const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; + const uint mbidxshift0 = (is < 4) ? 0 : 4; + const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + const uint mbidxshift1 = (is < 4) ? 0 : 2; + + const uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + const uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const float d = loadd.x * sc; + const float m = -loadd.y * mbyte; + + buf_a[buf_idx ] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi ] >> (b * 4)) & 0xF), m)); + buf_a[buf_idx + 1] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi + 1] >> (b * 4)) & 0xF), m)); +#elif defined(DATA_A_Q5_K) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint n = iqs / 32; // 0,1,2,3 + const uint b = (iqs % 32) / 16; // 0,1 + const uint is = 2 * n + b; // 0..7 + const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 + const uint qhi = (iqs % 16) * 2; // 0,2,4..30 + + const uint8_t hm = uint8_t(1 << (iqs / 16)); + + const vec2 loadd = vec2(data_a[ib].d); + + const uint scidx0 = (is < 4) ? is : (is + 4); + const uint scidx1 = (is < 4) ? is : (is - 4); + const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; + const uint scidxshift1 = (is < 4) ? 0 : 2; + const uint mbidx0 = is + 4; + const uint mbidx1 = (is < 4) ? is + 4 : is; + const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; + const uint mbidxshift0 = (is < 4) ? 0 : 4; + const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + const uint mbidxshift1 = (is < 4) ? 0 : 2; + + const uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + const uint8_t mbyte = uint8_t(((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const float d = loadd.x * sc; + const float m = -loadd.y * mbyte; + + buf_a[buf_idx ] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi ] >> (b * 4)) & 0xF) + float((data_a[ib].qh[qhi ] & hm) != 0 ? 16 : 0), m)); + buf_a[buf_idx + 1] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi + 1] >> (b * 4)) & 0xF) + float((data_a[ib].qh[qhi + 1] & hm) != 0 ? 16 : 0), m)); +#elif defined(DATA_A_Q6_K) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint n = iqs / 64; // 0,1 + const uint b = (iqs % 64) / 32; // 0,1 + const uint is_b = (iqs % 16) / 8; // 0,1 + const uint qhshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 + const uint is = 8 * n + qhshift + is_b; // 0..15 + const uint qsi = n * 64 + (iqs % 32) * 2; // 0,2,4..126 + const uint qhi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 + + const float dscale = float(data_a[ib].d) * float(data_a[ib].scales[is]); + + buf_a[buf_idx ] = FLOAT_TYPE(dscale * float(int8_t(((data_a[ib].ql[qsi ] >> (b * 4)) & 0xF) | (((data_a[ib].qh[qhi ] >> qhshift) & 3) << 4)) - 32)); + buf_a[buf_idx + 1] = FLOAT_TYPE(dscale * float(int8_t(((data_a[ib].ql[qsi + 1] >> (b * 4)) & 0xF) | (((data_a[ib].qh[qhi + 1] >> qhshift) & 3) << 4)) - 32)); +#elif defined(DATA_A_IQ4_NL) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; + + const uint ib = idx / 16; + const uint iqs = idx & 0xF; + + const float d = float(data_a[ib].d); + const uint vui = uint(data_a[ib].qs[iqs]); + const vec2 v = vec2(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[vui >> 4]) * d; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); +#endif + } + [[unroll]] for (uint l = 0; l < BN; l += loadstride_b) { +#if LOAD_VEC_B == 8 +#ifdef MUL_MAT_ID + const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l]; + const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + loadr_b; +#else + const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b; +#endif + const uint buf_idx = (loadc_b + l) * SHMEM_STRIDE + loadr_b * LOAD_VEC_B; + buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx][0].x); + buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx][0].y); + buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx][0].z); + buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx][0].w); + buf_b[buf_idx + 4] = FLOAT_TYPE(data_b[idx][1].x); + buf_b[buf_idx + 5] = FLOAT_TYPE(data_b[idx][1].y); + buf_b[buf_idx + 6] = FLOAT_TYPE(data_b[idx][1].z); + buf_b[buf_idx + 7] = FLOAT_TYPE(data_b[idx][1].w); +#elif LOAD_VEC_B == 4 +#ifdef MUL_MAT_ID + const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l]; + const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + loadr_b; +#else + const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b; +#endif + const uint buf_idx = (loadc_b + l) * SHMEM_STRIDE + loadr_b * LOAD_VEC_B; + buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx].x); + buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx].y); + buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx].z); + buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx].w); +#elif !MUL_MAT_ID + if (ic * BN + loadc_b + l < p.N && block + loadr_b < end_k) { + buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(data_b[pos_b + (loadc_b + l) * p.stride_b + loadr_b]); + } else { + buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(0.0f); + } +#else + const uint row_i = ic * BN + loadc_b + l; + if (row_i < _ne1) { + const u16vec2 row_idx = row_ids[row_i]; + buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(data_b[pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + loadr_b]); + } else { + buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(0.0f); + } +#endif + } + + barrier(); + + pos_a += BK / LOAD_VEC_A; + pos_b += BK / LOAD_VEC_B; + +#ifdef COOPMAT + [[unroll]] for (uint i = 0; i < BK; i += TK) { + [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { + // Load from shared into cache + coopMatLoad(cache_a, buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { + coopMatLoad(cache_b, buf_b, (warp_c * WN + cm_col * TN) * SHMEM_STRIDE + i, SHMEM_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + + sums[cm_col * cms_per_row + cm_row] = coopMatMulAdd(cache_a, cache_b, sums[cm_col * cms_per_row + cm_row]); + } + } + } +#else + [[unroll]] for (uint i = 0; i < BK; i++) { + // Load from shared into cache + [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { + [[unroll]] for (uint j = 0; j < TM; j++) { + cache_a[wsir * TM + j] = buf_a[(warp_r * WM + wsir * WSUBM + tiwr * TM + j) * SHMEM_STRIDE + i]; + } + } + [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { + [[unroll]] for (uint j = 0; j < TN; j++) { + cache_b[wsic * TN + j] = buf_b[(warp_c * WN + wsic * WSUBN + tiwc * TN + j) * SHMEM_STRIDE + i]; + } + } + + [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { + [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { + [[unroll]] for (uint cc = 0; cc < TN; cc++) { + [[unroll]] for (uint cr = 0; cr < TM; cr++) { + const uint sums_idx = (wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr; + sums[sums_idx] = fma(ACC_TYPE(cache_a[wsir * TM + cr]), ACC_TYPE(cache_b[wsic * TN + cc]), sums[sums_idx]); + } + } + } + } + } +#endif + + barrier(); + } + + const uint dr = ir * BM + warp_r * WM; + const uint dc = ic * BN + warp_c * WN; + +#ifndef MUL_MAT_ID + const uint offsets = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * gl_NumWorkGroups.z; +#endif + +#ifdef COOPMAT +#ifdef MUL_MAT_ID + [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { + [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { + coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); + + [[unroll]] for (uint col = 0; col < BN; col += storestride) { + const uint row_i = dc + cm_col * TN + col + store_c; + if (row_i >= _ne1) break; + + const u16vec2 row_idx = row_ids[row_i]; + + data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); + } + } + } +#else + const bool is_aligned = p.stride_d % 4 == 0; // Assumption: D_TYPE == float + + [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { + [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { + const bool is_in_bounds = dr + (cm_row + 1) * TM <= p.M && dc + (cm_col + 1) * TN <= p.N; + + if (is_aligned && is_in_bounds) { + // Full coopMat is within bounds and stride_d is aligned with 16B + coopmat cm_dtype = coopmat(sums[cm_col * cms_per_row + cm_row]); + coopMatStore(cm_dtype, data_d, offsets + (dc + cm_col * TN) * p.stride_d + dr + cm_row * TM, p.stride_d, gl_CooperativeMatrixLayoutColumnMajor); + } else if (is_in_bounds) { + // Full coopMat is within bounds, but stride_d is not aligned + coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); + + [[unroll]] for (uint col = 0; col < TN; col += storestride) { + data_d[offsets + (dc + cm_col * TN + col + store_c) * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); + } + } else if (dr + cm_row * TM < p.M && dc + cm_col * TN < p.N) { + // Partial coopMat is within bounds + coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); + + [[unroll]] for (uint col = 0; col < TN; col += storestride) { + if (dr + cm_row * TM + store_r < p.M && dc + cm_col * TN + col + store_c < p.N) { + data_d[offsets + (dc + cm_col * TN + col + store_c) * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); + } + } + } + } + } +#endif // MUL_MAT_ID +#else + [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { + [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { + + const uint dr_warp = dr + wsir * WSUBM + tiwr * TM; + const uint dc_warp = dc + wsic * WSUBN + tiwc * TN; + [[unroll]] for (uint cc = 0; cc < TN; cc++) { +#ifdef MUL_MAT_ID + const uint row_i = dc_warp + cc; + if (row_i >= _ne1) break; + + const u16vec2 row_idx = row_ids[row_i]; +#endif // MUL_MAT_ID + [[unroll]] for (uint cr = 0; cr < TM; cr++) { +#ifdef MUL_MAT_ID + data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]); +#else + if (dr_warp + cr < p.M && dc_warp + cc < p.N) { + data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]); + } +#endif // MUL_MAT_ID + } + } + } + } +#endif // COOPMAT +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp new file mode 100644 index 000000000..cbfa5dce1 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp @@ -0,0 +1,328 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require + +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +#extension GL_KHR_memory_scope_semantics : enable +#extension GL_KHR_cooperative_matrix : enable +#extension GL_NV_cooperative_matrix2 : enable +#extension GL_EXT_buffer_reference : enable +#extension GL_KHR_shader_subgroup_ballot : enable +#extension GL_KHR_shader_subgroup_vote : enable + +#include "types.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (constant_id = 1) const uint BM = 64; +layout (constant_id = 2) const uint BN = 64; +layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working with a quant + +layout (push_constant) uniform parameter +{ + uint M; + uint N; + uint K; + uint stride_a; + uint stride_b; + uint stride_d; + + uint batch_stride_a; + uint batch_stride_b; + uint batch_stride_d; + +#ifdef MUL_MAT_ID + uint nei0; + uint nei1; + uint nbi1; + uint ne11; +#else + uint k_split; + uint ne02; + uint ne12; + uint broadcast2; + uint broadcast3; +#endif +} p; + + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; + +#if QUANT_K > 1 +#define DECODEFUNCA , dequantFuncA +#define MAT_A_TYPE float16_t + +#include "dequant_funcs_cm2.comp" + +#else +#define DECODEFUNCA +#define MAT_A_TYPE A_TYPE +#endif + +#define MAT_B_TYPE B_TYPE + +#ifdef MUL_MAT_ID +layout (binding = 3) readonly buffer IDS {int data_ids[];}; + +shared u16vec4 row_ids[3072]; + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufB { + B_TYPE b[]; +}; + +uint _ne1; +shared uint _ne1_sh; + +B_TYPE decodeFuncB(const in decodeBufB bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint row_i = blockCoords[0]; + + if (row_i >= _ne1) { + return B_TYPE(0.0); + } + + const u16vec4 row_idx = row_ids[row_i]; + B_TYPE ret = data_b[row_idx.y * p.batch_stride_b + row_idx.x * p.stride_b + blockCoords[1]]; + + return ret; +} + +D_TYPE perElemOpD(const in uint32_t r, const in uint32_t c, const in D_TYPE elem, const in uint32_t ir, const in uint32_t ic) +{ + uint dr = ir * BM + r; + uint dc = ic * BN + c; + + if (dr < p.M && dc < _ne1) { + uint row_i = dc; + const u16vec4 row_idx = row_ids[row_i]; + data_d[row_idx.y * p.batch_stride_d + row_idx.z * p.stride_d + dr] = elem; + } + return elem; +} + +#endif + +void main() { +#if defined(DATA_A_IQ4_NL) + init_iq4nl_shmem(); +#endif + +#ifdef MUL_MAT_ID + const uint expert_idx = gl_GlobalInvocationID.z; +#else + const uint batch_idx = gl_GlobalInvocationID.z; + + const uint i13 = batch_idx / p.ne12; + const uint i12 = batch_idx % p.ne12; + + const uint i03 = i13 / p.broadcast3; + const uint i02 = i12 / p.broadcast2; + + const uint batch_idx_a = i03 * p.ne02 + i02; +#endif + + const uint blocks_m = (p.M + BM - 1) / BM; + const uint ir = gl_WorkGroupID.x % blocks_m; + const uint ik = gl_WorkGroupID.x / blocks_m; + const uint ic = gl_WorkGroupID.y; + +#ifdef MUL_MAT_ID + // Spread the search across all elements in the first subgroup + if (gl_SubgroupID == 0) { + _ne1 = 0; + uint num_elements = p.nei1 * p.nei0; + + for (uint i = gl_SubgroupInvocationID; subgroupAny(i < num_elements); i += gl_SubgroupSize) { + bool in_range = i < num_elements; + uint ii0 = i % p.nei0; + uint ii1 = i / p.nei0; + uint id = in_range ? data_ids[ii1*p.nbi1 + ii0] : 0; + uvec4 ballot = subgroupBallot(in_range && id == expert_idx); + uint idx = subgroupBallotExclusiveBitCount(ballot); + if (in_range && id == expert_idx) { + row_ids[_ne1 + idx] = u16vec4(ii0 % p.ne11, ii1, ii0, 0); + } + _ne1 += subgroupBallotBitCount(ballot); + } + _ne1_sh = _ne1; + } + + barrier(); + + _ne1 = _ne1_sh; + + // Workgroup has no work + if (ic * BN >= _ne1) return; +#endif + +#ifdef MUL_MAT_ID + uint start_k = 0; + const uint end_k = p.K; +#else + uint start_k = ik * p.k_split; + const uint end_k = min(p.K, (ik + 1) * p.k_split); +#endif + + coopmat sum; + sum = coopmat(0.0); + +#ifdef MUL_MAT_ID + uint pos_a = (expert_idx * p.batch_stride_a) / QUANT_K; + uint pos_b = 0; +#else + uint pos_a = (batch_idx_a * p.batch_stride_a) / QUANT_K; + uint pos_b = batch_idx * p.batch_stride_b; +#endif + + uint stride_a = p.stride_a / QUANT_K; + uint stride_b = p.stride_b; + + // Hint to the compiler that values are aligned (want 16B alignment). + // Quants are always block-aligned, no alignment needed. +#if ALIGNED +#if QUANT_K == 1 + stride_a &= ~7; +#endif + stride_b &= ~7; +#endif + + // Create layouts for both clamped and unclamped accesses + tensorLayoutNV<2> tensorLayoutA = createTensorLayoutNV(2); + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutAClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutNV<2> tensorLayoutB = createTensorLayoutNV(2); + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutBClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + +#if QUANT_K > 1 + tensorLayoutA = setTensorLayoutBlockSizeNV(tensorLayoutA, 1, QUANT_K); + tensorLayoutAClamp = setTensorLayoutBlockSizeNV(tensorLayoutAClamp, 1, QUANT_K); +#endif + + // Use end_k rather than p.K as the dimension because that's what + // we need to bound check against when using split_k + tensorLayoutA = setTensorLayoutDimensionNV(tensorLayoutA, p.M, end_k); + tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, p.N, end_k); + tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.N, p.M); + tensorLayoutAClamp = setTensorLayoutDimensionNV(tensorLayoutAClamp, p.M, end_k); + tensorLayoutBClamp = setTensorLayoutDimensionNV(tensorLayoutBClamp, p.N, end_k); + + tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0); + +#if !defined(MUL_MAT_ID) + // Detect a fast path where all loads are entirely in bounds and no clamping is required + if ((ir + 1) * BM <= p.M && (ic + 1) * BN <= p.N && (start_k % BK) == 0 && (end_k % BK) == 0 && +#if QUANT_K == 1 + (stride_a % 8) == 0 && +#endif + (stride_b % 8) == 0 && (start_k % 8) == 0) { + // Hint to the compiler that values are aligned (want 16B alignment) + start_k &= ~7; + stride_b &= ~7; +#if QUANT_K == 1 + stride_a &= ~7; +#endif + + tensorLayoutA = setTensorLayoutStrideNV(tensorLayoutA, stride_a, 1); + tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1); + + uint k_iters = (end_k - start_k + BK - 1) / BK; + + for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) { + + coopmat mat_a; + coopmat mat_b; + + coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA); + coopmat mat_a_ft = coopmat(mat_a); + + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose); + coopmat mat_b_ft = coopmat(mat_b); + + sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); + } + } else +#endif // !defined(MUL_MAT_ID) + { + tensorLayoutA = setTensorLayoutStrideNV(tensorLayoutA, stride_a, 1); + + tensorLayoutAClamp = setTensorLayoutStrideNV(tensorLayoutAClamp, stride_a, 1); + + tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1); + + tensorLayoutBClamp = setTensorLayoutStrideNV(tensorLayoutBClamp, stride_b, 1); + + [[dont_unroll]] + for (uint block_k = start_k; block_k < end_k; block_k += BK) { + + coopmat mat_a; + coopmat mat_b; + coopmat mat_a_ft; + coopmat mat_b_ft; + + // Clamping is expensive, so detect different code paths for each combination + // of A and B needing clamping. + bool unclampedA = (ir + 1) * BM <= p.M && block_k + BK <= end_k && (block_k % 8) == 0; +#ifdef MUL_MAT_ID + bool unclampedB = true; +#else + bool unclampedB = (ic + 1) * BN <= p.N && block_k + BK <= end_k && (block_k % 8) == 0; +#endif + if (unclampedA && unclampedB) { + coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, (block_k & ~7), BK) DECODEFUNCA); +#ifdef MUL_MAT_ID + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose, decodeFuncB); +#else + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, (block_k & ~7), BK), tensorViewTranspose); +#endif + mat_a_ft = coopmat(mat_a); + mat_b_ft = coopmat(mat_b); + sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); + } else if (unclampedA && !unclampedB) { + coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, (block_k & ~7), BK) DECODEFUNCA); + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutBClamp, ic * BN, BN, block_k, BK), tensorViewTranspose); + + mat_a_ft = coopmat(mat_a); + mat_b_ft = coopmat(mat_b); + sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); + } else if (!unclampedA && unclampedB) { + coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutAClamp, ir * BM, BM, block_k, BK) DECODEFUNCA); +#ifdef MUL_MAT_ID + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose, decodeFuncB); +#else + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, (block_k & ~7), BK), tensorViewTranspose); +#endif + mat_a_ft = coopmat(mat_a); + mat_b_ft = coopmat(mat_b); + sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); + } else if (!unclampedA && !unclampedB) { + coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutAClamp, ir * BM, BM, block_k, BK) DECODEFUNCA); + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutBClamp, ic * BN, BN, block_k, BK), tensorViewTranspose); + + mat_a_ft = coopmat(mat_a); + mat_b_ft = coopmat(mat_b); + sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); + } + } + } + + // Convert from ACC_TYPE to D_TYPE + coopmat mat_d; + mat_d = coopmat(sum); + +#ifdef MUL_MAT_ID + // Call callback to store each element, remapping row through shared memory + coopMatPerElementNV(mat_d, mat_d, perElemOpD, ir, ic); +#else + tensorLayoutD = setTensorLayoutStrideNV(tensorLayoutD, p.stride_d, 1); + + uint pos_d = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * gl_NumWorkGroups.z; + coopMatStoreTensorNV(mat_d, data_d, pos_d, sliceTensorLayoutNV(tensorLayoutD, ic * BN, BN, ir * BM, BM), tensorViewTranspose); +#endif +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp new file mode 100644 index 000000000..6627a50bd --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp @@ -0,0 +1,44 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable +#define BLOCK_SIZE 512 + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +shared vec2 sum[BLOCK_SIZE]; + +void main() { + const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationID.x; + + sum[tid] = vec2(0.0f, 0.0f); + + [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { + const float xi = float(data_a[row*p.KX + col]); + sum[tid].x += xi; + sum[tid].y += xi * xi; + } + + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + sum[tid] += sum[tid + s]; + } + barrier(); + } + + const float mean = sum[0].x / p.KX; + const float var = sum[0].y / p.KX - mean * mean; + const float inv_std = inversesqrt(var + p.param1); + + [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { + data_d[row*p.KX + col] = D_TYPE((float(data_a[row*p.KX + col]) - mean) * inv_std); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp new file mode 100644 index 000000000..450b67fc5 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp @@ -0,0 +1,28 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (idx >= p.ne) { + return; + } + + const uint i3 = idx / (p.ne12*p.ne11*p.ne10); + const uint i3_offset = i3 * p.ne12*p.ne11*p.ne10; + const uint i2 = (idx - i3_offset) / (p.ne11*p.ne10); + const uint i2_offset = i2*p.ne11*p.ne10; + const uint i1 = (idx - i3_offset - i2_offset) / p.ne10; + const uint i0 = idx - i3_offset - i2_offset - i1*p.ne10; + + const uint src0_idx = i3*p.nb03 + i2*p.nb02 + i1*p.nb01 + i0*p.nb00; + const uint dst_idx = i3*p.nb13 + i2*p.nb12 + i1*p.nb11 + i0*p.nb10; + + const bool is_src0 = i0 < p.ne00 && i1 < p.ne01 && i2 < p.ne02 && i3 < p.ne03; + + data_d[get_doffset() + dst_idx] = D_TYPE(is_src0 ? data_a[get_aoffset() + src0_idx] : 0.0f); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp new file mode 100644 index 000000000..b6124411a --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp @@ -0,0 +1,74 @@ +#version 450 + +#include "types.comp" + +#extension GL_EXT_shader_16bit_storage : require + +layout(push_constant) uniform parameter { + uint IW; uint IH; + uint OW; uint OH; + uint OC; + uint pelements; + uint op; + int k0; int k1; + int s0; int s1; + int p0; int p1; +} p; + +#define BLOCK_SIZE 512 +#define FLT_MAX 3.402823466e+38F +#define OP_POOL_MAX 0u +#define OP_POOL_AVG 1u + +layout (local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout(binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint idx = gl_GlobalInvocationID.x; + if (idx >= p.pelements) { + return; + } + + const uint O_HW = p.OW * p.OH; + + const uint nc = idx / O_HW; + const uint cur_oh = (idx % O_HW) / p.OW; + const uint cur_ow = (idx % O_HW) % p.OW; + + const int start_h = int(cur_oh) * p.s0 - p.p0; + const uint bh = max(start_h, 0); + const uint eh = min(start_h + p.k0, p.IH); + + const int start_w = int(cur_ow) * p.s1 - p.p1; + const uint bw = max(start_w, 0); + const uint ew = min(start_w + p.k1, p.IW); + + const float scale = 1.0 / float(p.k0 * p.k1); + float res; + + if (p.op == OP_POOL_AVG) { + res = 0.0; + } else if (p.op == OP_POOL_MAX) { + res = -FLT_MAX; + } else { + return; + } + + #pragma unroll + for (uint i = bh; i < eh; i++) { + #pragma unroll + for (uint j = bw; j < ew; j++) { + const float cur = D_TYPE(data_a[nc * p.IH * p.IW + i * p.IW + j]); + + if (p.op == OP_POOL_AVG) { + res += cur * scale; + } else if (p.op == OP_POOL_MAX) { + res = max(res, cur); + } + } + } + + data_d[nc * O_HW + cur_oh * p.OW + cur_ow] = res; +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp new file mode 100644 index 000000000..52a19b62a --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp @@ -0,0 +1,21 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + + data_d[i] = max(float(data_a[i]), 0); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp new file mode 100644 index 000000000..1568b141d --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp @@ -0,0 +1,26 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +uint src0_idx_mod(uint idx) { + const uint i13 = idx / (p.ne12*p.ne11*p.ne10); + const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10; + const uint i12 = (idx - i13_offset) / (p.ne11*p.ne10); + const uint i12_offset = i12*p.ne11*p.ne10; + const uint i11 = (idx - i13_offset - i12_offset) / p.ne10; + const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10; + return (i13 % p.ne03)*p.nb03 + (i12 % p.ne02)*p.nb02 + (i11 % p.ne01)*p.nb01 + (i10 % p.ne00)*p.nb00; +} + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(data_a[get_aoffset() + src0_idx_mod(idx)]); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp new file mode 100644 index 000000000..b554400ba --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp @@ -0,0 +1,42 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable +#define BLOCK_SIZE 512 + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +shared FLOAT_TYPE sum[BLOCK_SIZE]; + +void main() { + const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationID.x; + + sum[tid] = FLOAT_TYPE(0.0f); // partial sum for thread in warp + + [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { + const FLOAT_TYPE xi = FLOAT_TYPE(data_a[row*p.KX + col]); + sum[tid] += xi * xi; + } + + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + sum[tid] += sum[tid + s]; + } + barrier(); + } + + const FLOAT_TYPE mean = sum[0] / FLOAT_TYPE(p.KX); + const FLOAT_TYPE scale = inversesqrt(mean + FLOAT_TYPE(p.param1)); + + [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { + data_d[row*p.KX + col] = D_TYPE(scale * FLOAT_TYPE(data_a[row*p.KX + col])); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp new file mode 100644 index 000000000..574b51ca5 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp @@ -0,0 +1,49 @@ +#include "types.comp" + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_spirv_intrinsics: enable + +#if RTE16 +spirv_execution_mode(capabilities = [4467], 4462, 16); // RoundingModeRTE, 16 bits +#endif + +layout(local_size_x = 1, local_size_y = 256, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer Y {int data_pos[];}; +layout (binding = 2) readonly buffer Z {float data_ff[];}; +layout (binding = 3) writeonly buffer D {D_TYPE data_d[];}; + +layout (push_constant) uniform parameter { + uint ncols; + uint n_dims; + float freq_scale; + uint p_delta_rows; + float freq_base; + float ext_factor; + float attn_factor; + float corr_dims[2]; + float theta_scale; + uint has_ff; +} p; + +float rope_yarn_ramp(const float low, const float high, const uint i0) { + const float y = (i0 / 2 - low) / max(0.001f, high - low); + return 1.0f - min(1.0f, max(0.0f, y)); +} + +void rope_yarn(const float theta_extrap, const uint i0, out float cos_theta, out float sin_theta) { + float mscale = p.attn_factor; + // Get n-d rotational scaling corrected for extrapolation + float theta_interp = p.freq_scale * theta_extrap; + float theta = theta_interp; + if (p.ext_factor != 0.0f) { + float ramp_mix = rope_yarn_ramp(p.corr_dims[0], p.corr_dims[1], i0) * p.ext_factor; + theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix; + + // Get n-d magnitude scaling corrected for interpolation + mscale *= 1.0f + 0.1f * log(1.0f / p.freq_scale); + } + cos_theta = cos(theta) * mscale; + sin_theta = sin(theta) * mscale; +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp new file mode 100644 index 000000000..83b46b69b --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp @@ -0,0 +1,37 @@ +#version 450 + +#include "rope_head.comp" + +void main() { + const uint col = gl_GlobalInvocationID.y * 2; + const uint row = gl_GlobalInvocationID.x; + + if (col >= p.ncols) { + return; + } + + if (col >= p.n_dims) { + const uint i = row*p.ncols + col; + + data_d[i + 0] = data_a[i + 0]; + data_d[i + 1] = data_a[i + 1]; + + return; + } + + const uint i = row*p.ncols + col/2; + const uint i2 = row/p.p_delta_rows; + + const float theta_base = data_pos[i2] * pow(p.theta_scale, col/2.0f); + + const float freq_factor = p.has_ff != 0 ? data_ff[col/2] : 1.0f; + + float cos_theta, sin_theta; + rope_yarn(theta_base / freq_factor, col, cos_theta, sin_theta); + + const float x0 = float(data_a[i + 0]); + const float x1 = float(data_a[i + p.n_dims/2]); + + data_d[i + 0] = D_TYPE(x0*cos_theta - x1*sin_theta); + data_d[i + p.n_dims/2] = D_TYPE(x0*sin_theta + x1*cos_theta); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp new file mode 100644 index 000000000..e416ad938 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp @@ -0,0 +1,37 @@ +#version 450 + +#include "rope_head.comp" + +void main() { + const uint col = gl_GlobalInvocationID.y * 2; + const uint row = gl_GlobalInvocationID.x; + + if (col >= p.ncols) { + return; + } + + if (col >= p.n_dims) { + const uint i = row*p.ncols + col; + + data_d[i + 0] = data_a[i + 0]; + data_d[i + 1] = data_a[i + 1]; + + return; + } + + const uint i = row*p.ncols + col; + const uint i2 = row/p.p_delta_rows; + + const float theta_base = data_pos[i2] * pow(p.theta_scale, col/2.0f); + + const float freq_factor = p.has_ff != 0 ? data_ff[col/2] : 1.0f; + + float cos_theta, sin_theta; + rope_yarn(theta_base / freq_factor, col, cos_theta, sin_theta); + + const float x0 = float(data_a[i + 0]); + const float x1 = float(data_a[i + 1]); + + data_d[i + 0] = D_TYPE(x0*cos_theta - x1*sin_theta); + data_d[i + 1] = D_TYPE(x0*sin_theta + x1*cos_theta); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp new file mode 100644 index 000000000..4663428de --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp @@ -0,0 +1,24 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +const uint num_threads = 128; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + uint idx = get_idx(); + + // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation + const uint num_iter = 4; + + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (idx >= p.ne) { + continue; + } + + data_d[get_doffset() + idx] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + idx]) * FLOAT_TYPE(p.param1)); + idx += num_threads; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp new file mode 100644 index 000000000..4d36f88e0 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp @@ -0,0 +1,22 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + + const float xi = float(data_a[i]); + data_d[i] = D_TYPE(xi / (1.0f + exp(-xi))); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp new file mode 100644 index 000000000..d7c15a169 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp @@ -0,0 +1,17 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + + const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(sin(val)); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp new file mode 100644 index 000000000..a25808e16 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp @@ -0,0 +1,174 @@ +#version 450 + +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +layout (push_constant) uniform parameter +{ + uint KX; + uint KY; + float scale; + float max_bias; + float m0; + float m1; + uint n_head_log2; + uint nrows_x; +} p; + +#include "types.comp" + +layout(constant_id = 0) const uint BLOCK_SIZE = 32; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer Y {B_TYPE data_b[];}; +layout (binding = 2) buffer D {D_TYPE data_d[];}; + +shared FLOAT_TYPE vals[BLOCK_SIZE]; + +// num_iters is the number of BLOCK_SIZE loop iterations we need to iterate +// over all the columns. The main function tries to pass a constant here, +// as if it were a template function, to allow unrolling. +void soft_max(uint num_iters) { + const uint tid = gl_LocalInvocationID.x; + const uint rowx = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; + const uint rowy = (p.KY > 0) ? (rowx % p.KY) : 0; + + if (rowx >= p.nrows_x) { + return; + } + + float slope = 1.0f; + + // ALiBi + if (p.max_bias > 0.0f) { + const uint h = rowx/p.KY; // head index + + const float base = h < p.n_head_log2 ? p.m0 : p.m1; + const uint exp = h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1; + + slope = pow(base, exp); + } + + // Find max + FLOAT_TYPE max_val = uintBitsToFloat(0xFF800000); + + // Cache values while we compute the max, so we don't need to read them + // again when we're ready to compute exp(x-max). + const uint DATA_CACHE_SIZE = 16; + FLOAT_TYPE data_cache[DATA_CACHE_SIZE]; + + [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { + const uint col = col0 + tid; + + FLOAT_TYPE a = FLOAT_TYPE(0); + if (col < p.KX) { + a = data_a[rowx * p.KX + col]; + } + + FLOAT_TYPE b = FLOAT_TYPE(0); + if (p.KY > 0 && col < p.KX) { + b = data_b[rowy * p.KX + col]; + } + + FLOAT_TYPE v = a * p.scale + slope * b; + + if (col < p.KX) { + max_val = max(max_val, v); + } + + if (idx < DATA_CACHE_SIZE) { + data_cache[idx] = v; + } + } + + // reduce across the workgroup + vals[tid] = max_val; + barrier(); + [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + vals[tid] = max(vals[tid], vals[tid + s]); + } + barrier(); + } + + max_val = vals[0]; + barrier(); + + FLOAT_TYPE sum = FLOAT_TYPE(0.0f); + + // Compute sum{exp(x - max)} + [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { + const uint col = col0 + tid; + + if (col >= p.KX) { + break; + } + + // compute exp(a*scale+b*slope), add it to sum, and cache the new value + // in data_cache if possible. + const uint i = rowx * p.KX + col; + FLOAT_TYPE val; + if (idx < DATA_CACHE_SIZE) { + val = exp(data_cache[idx] - max_val); + } else { + val = exp(FLOAT_TYPE(data_a[i]) * p.scale + (p.KY > 0 ? slope * FLOAT_TYPE(data_b[rowy * p.KX + col]) : FLOAT_TYPE(0.0f)) - max_val); + } + sum += val; + if (idx < DATA_CACHE_SIZE) { + data_cache[idx] = val; + } else { + data_d[i] = D_TYPE(val); + } + } + + // reduce across the workgroup + vals[tid] = sum; + barrier(); + [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + vals[tid] += vals[tid + s]; + } + barrier(); + } + sum = vals[0]; + + FLOAT_TYPE rcpdivisor = 1.0/sum; + + [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { + const uint col = col0 + tid; + + if (col >= p.KX) { + continue; + } + + if (idx < DATA_CACHE_SIZE) { + data_d[rowx*p.KX + col] = D_TYPE(data_cache[idx] * rcpdivisor); + } else { + data_d[rowx*p.KX + col] *= D_TYPE(rcpdivisor); + } + } +} + +void main() { + // instantiate the soft_max function for several different + // dimensions, to allow loop unrolling + uint num_blocks = (p.KX + BLOCK_SIZE - 1) / BLOCK_SIZE; + if (num_blocks > 32) { + soft_max(num_blocks); + } else if (num_blocks > 16) { + soft_max(32); + } else if (num_blocks > 8) { + soft_max(16); + } else if (num_blocks > 4) { + soft_max(8); + } else if (num_blocks == 4) { + soft_max(4); + } else if (num_blocks == 3) { + soft_max(3); + } else if (num_blocks == 2) { + soft_max(2); + } else if (num_blocks == 1) { + soft_max(1); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp new file mode 100644 index 000000000..ef43598ba --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp @@ -0,0 +1,17 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + + const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(val * val); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp new file mode 100644 index 000000000..961e5ffa1 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp @@ -0,0 +1,37 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +layout (constant_id = 0) const uint BLOCK_SIZE = 32; + +shared FLOAT_TYPE tmp[BLOCK_SIZE]; + +void main() { + const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; + const uint col = gl_LocalInvocationID.x; + + tmp[col] = FLOAT_TYPE(0.0f); + + for (uint i = col; i < p.KX; i += BLOCK_SIZE) { + tmp[col] += FLOAT_TYPE(data_a[row*p.KX + i]); + } + + barrier(); + [[unroll]] for (int s = int(BLOCK_SIZE) / 2; s > 0; s >>= 1) { + if (col < s) { + tmp[col] += tmp[col + s]; + } + barrier(); + } + + if (col == 0) { + data_d[row] = D_TYPE(tmp[0]); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp new file mode 100644 index 000000000..495f966bd --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp @@ -0,0 +1,20 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + data_d[i] = D_TYPE(1. - 2. / (exp(2.*data_a[i]) + 1.)); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp new file mode 100644 index 000000000..28eb24e11 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp @@ -0,0 +1,7 @@ +#version 460 + +#extension GL_NV_cooperative_matrix2 : require + +void main() +{ +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp new file mode 100644 index 000000000..79e065a93 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp @@ -0,0 +1,41 @@ +#version 450 + +#extension GL_EXT_shader_16bit_storage : require + +layout (push_constant) uniform parameter +{ + uint nb1; + uint dim; + uint max_period; +} p; + +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable +#define BLOCK_SIZE 256 + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint i = gl_WorkGroupID.y; + const uint j = gl_GlobalInvocationID.x; + const uint d_offset = i * p.nb1; + + if (p.dim % 2 != 0 && j == ((p.dim + 1) / 2)) { + data_d[d_offset + p.dim] = 0.f; + } + + const uint half_dim = p.dim / 2; + if (j >= half_dim) { + return; + } + + const float timestep = float(data_a[i]); + const float freq = float(exp(-log(p.max_period) * j / half_dim)); + const float arg = timestep * freq; + data_d[d_offset + j] = D_TYPE(cos(arg)); + data_d[d_offset + j + half_dim] = D_TYPE(sin(arg)); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp new file mode 100644 index 000000000..eecc47f3a --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp @@ -0,0 +1,323 @@ + +#if !defined(GGML_TYPES_COMP) +#define GGML_TYPES_COMP + +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#if defined(DATA_A_F32) +#define QUANT_K 1 +#define QUANT_R 1 + +#if !defined(LOAD_VEC_A) || LOAD_VEC_A == 1 +#define A_TYPE float +#elif LOAD_VEC_A == 4 +#define A_TYPE vec4 +#elif LOAD_VEC_A == 8 +#define A_TYPE mat2x4 +#endif +#endif + +#if defined(DATA_A_F16) +#define QUANT_K 1 +#define QUANT_R 1 + +#if !defined(LOAD_VEC_A) || LOAD_VEC_A == 1 +#define A_TYPE float16_t +#elif LOAD_VEC_A == 4 +#define A_TYPE f16vec4 +#elif LOAD_VEC_A == 8 +#define A_TYPE f16mat2x4 +#endif +#endif + +#define QUANT_K_Q4_0 32 +#define QUANT_R_Q4_0 2 + +struct block_q4_0 +{ + float16_t d; + uint8_t qs[16]; +}; +struct block_q4_0_packed16 +{ + float16_t d; + uint16_t qs[16/2]; +}; + +#if defined(DATA_A_Q4_0) +#define QUANT_K QUANT_K_Q4_0 +#define QUANT_R QUANT_R_Q4_0 +#define A_TYPE block_q4_0 +#define A_TYPE_PACKED16 block_q4_0_packed16 +#endif + +#define QUANT_K_Q4_1 32 +#define QUANT_R_Q4_1 2 + +struct block_q4_1 +{ + float16_t d; + float16_t m; + uint8_t qs[16]; +}; + +struct block_q4_1_packed16 +{ + float16_t d; + float16_t m; + uint16_t qs[16/2]; +}; + +#if defined(DATA_A_Q4_1) +#define QUANT_K QUANT_K_Q4_1 +#define QUANT_R QUANT_R_Q4_1 +#define A_TYPE block_q4_1 +#define A_TYPE_PACKED16 block_q4_1_packed16 +#endif + +#define QUANT_K_Q5_0 32 +#define QUANT_R_Q5_0 2 + +struct block_q5_0 +{ + float16_t d; + uint16_t qh[2]; + uint8_t qs[16]; +}; + +struct block_q5_0_packed16 +{ + float16_t d; + uint16_t qh[2]; + uint16_t qs[16/2]; +}; + +#if defined(DATA_A_Q5_0) +#define QUANT_K QUANT_K_Q5_0 +#define QUANT_R QUANT_R_Q5_0 +#define A_TYPE block_q5_0 +#define A_TYPE_PACKED16 block_q5_0_packed16 +#endif + +#define QUANT_K_Q5_1 32 +#define QUANT_R_Q5_1 2 + +struct block_q5_1 +{ + float16_t d; + float16_t m; + uint qh; + uint8_t qs[16]; +}; + +struct block_q5_1_packed16 +{ + float16_t d; + float16_t m; + uint qh; + uint16_t qs[16/2]; +}; + +#if defined(DATA_A_Q5_1) +#define QUANT_K QUANT_K_Q5_1 +#define QUANT_R QUANT_R_Q5_1 +#define A_TYPE block_q5_1 +#define A_TYPE_PACKED16 block_q5_1_packed16 +#endif + +#define QUANT_K_Q8_0 32 +#define QUANT_R_Q8_0 1 + +struct block_q8_0 +{ + float16_t d; + int8_t qs[32]; +}; +struct block_q8_0_packed16 +{ + float16_t d; + uint16_t qs[32/2]; +}; + +#if defined(DATA_A_Q8_0) +#define QUANT_K QUANT_K_Q8_0 +#define QUANT_R QUANT_R_Q8_0 +#define A_TYPE block_q8_0 +#define A_TYPE_PACKED16 block_q8_0_packed16 +#endif + +// K-quants +#define QUANT_K_Q2_K 256 + +struct block_q2_K +{ + uint8_t scales[QUANT_K_Q2_K/16]; + uint8_t qs[QUANT_K_Q2_K/4]; + f16vec2 d; +}; + +struct block_q2_K_packed16 +{ + uint16_t scales[QUANT_K_Q2_K/16/2]; + uint16_t qs[QUANT_K_Q2_K/4/2]; + f16vec2 d; +}; + +struct block_q2_K_packed32 +{ + uint32_t scales[QUANT_K_Q2_K/16/4]; + uint32_t qs[QUANT_K_Q2_K/4/4]; + f16vec2 d; +}; + +#if defined(DATA_A_Q2_K) +#define QUANT_K QUANT_K_Q2_K +#define A_TYPE block_q2_K +#define A_TYPE_PACKED16 block_q2_K_packed16 +#define A_TYPE_PACKED32 block_q2_K_packed32 +#endif + +#define QUANT_K_Q3_K 256 + +struct block_q3_K +{ + uint8_t hmask[QUANT_K_Q3_K/8]; + uint8_t qs[QUANT_K_Q3_K/4]; + uint8_t scales[12]; + float16_t d; +}; + +struct block_q3_K_packed16 +{ + uint16_t hmask[QUANT_K_Q3_K/8/2]; + uint16_t qs[QUANT_K_Q3_K/4/2]; + uint16_t scales[12/2]; + float16_t d; +}; + +#if defined(DATA_A_Q3_K) +#define QUANT_K QUANT_K_Q3_K +#define A_TYPE block_q3_K +#define A_TYPE_PACKED16 block_q3_K_packed16 +#endif + +#define QUANT_K_Q4_K 256 + +struct block_q4_K +{ + f16vec2 d; + uint8_t scales[3*QUANT_K_Q4_K/64]; + uint8_t qs[QUANT_K_Q4_K/2]; +}; + +struct block_q4_K_packed16 +{ + f16vec2 d; + uint16_t scales[3*QUANT_K_Q4_K/64/2]; + uint16_t qs[QUANT_K_Q4_K/2/2]; +}; + +struct block_q4_K_packed32 +{ + f16vec2 d; + uint32_t scales[3*QUANT_K_Q4_K/64/4]; + uint32_t qs[QUANT_K_Q4_K/2/4]; +}; + +#if defined(DATA_A_Q4_K) +#define QUANT_K QUANT_K_Q4_K +#define A_TYPE block_q4_K +#define A_TYPE_PACKED16 block_q4_K_packed16 +#define A_TYPE_PACKED32 block_q4_K_packed32 +#endif + +#define QUANT_K_Q5_K 256 + +struct block_q5_K +{ + f16vec2 d; + uint8_t scales[12]; + uint8_t qh[QUANT_K_Q5_K/8]; + uint8_t qs[QUANT_K_Q5_K/2]; +}; + +struct block_q5_K_packed16 +{ + f16vec2 d; + uint16_t scales[12/2]; + uint16_t qh[QUANT_K_Q5_K/8/2]; + uint16_t qs[QUANT_K_Q5_K/2/2]; +}; + +#if defined(DATA_A_Q5_K) +#define QUANT_K QUANT_K_Q5_K +#define A_TYPE block_q5_K +#define A_TYPE_PACKED16 block_q5_K_packed16 +#endif + +#define QUANT_K_Q6_K 256 + +struct block_q6_K +{ + uint8_t ql[QUANT_K_Q6_K/2]; + uint8_t qh[QUANT_K_Q6_K/4]; + int8_t scales[QUANT_K_Q6_K/16]; + float16_t d; +}; + +struct block_q6_K_packed16 +{ + uint16_t ql[QUANT_K_Q6_K/2/2]; + uint16_t qh[QUANT_K_Q6_K/4/2]; + int8_t scales[QUANT_K_Q6_K/16]; + float16_t d; +}; + +#if defined(DATA_A_Q6_K) +#define QUANT_K QUANT_K_Q6_K +#define A_TYPE block_q6_K +#define A_TYPE_PACKED16 block_q6_K_packed16 +#endif + +// IQuants + +#define QUANT_K_IQ4_NL 32 +#define QUANT_R_IQ4_NL 2 + +struct block_iq4_nl +{ + float16_t d; + uint8_t qs[QUANT_K_IQ4_NL/2]; +}; + +struct block_iq4_nl_packed16 +{ + float16_t d; + uint16_t qs[QUANT_K_IQ4_NL/2/2]; +}; + +#if defined(DATA_A_IQ4_NL) + +const int8_t kvalues_iq4nl_const[16] = { + int8_t(-127), int8_t(-104), int8_t(-83), int8_t(-65), int8_t(-49), int8_t(-35), int8_t(-22), int8_t(-10), + int8_t(1), int8_t(13), int8_t(25), int8_t(38), int8_t(53), int8_t(69), int8_t(89), int8_t(113) +}; + +shared FLOAT_TYPE kvalues_iq4nl[16]; + +void init_iq4nl_shmem() +{ + // copy the table into shared memory and sync + if (gl_LocalInvocationIndex.x < 16) { + kvalues_iq4nl[gl_LocalInvocationIndex.x] = FLOAT_TYPE(kvalues_iq4nl_const[gl_LocalInvocationIndex.x]); + } + barrier(); +} + +#define QUANT_K QUANT_K_IQ4_NL +#define QUANT_R QUANT_R_IQ4_NL +#define A_TYPE block_iq4_nl +#define A_TYPE_PACKED16 block_iq4_nl_packed16 +#endif + +#endif // !defined(GGML_TYPES_COMP) diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp new file mode 100644 index 000000000..6f607380d --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp @@ -0,0 +1,36 @@ +#version 450 + +layout (push_constant) uniform parameter +{ + uint ne; uint a_offset; uint d_offset; + uint nb00; uint nb01; uint nb02; uint nb03; + uint ne10; uint ne11; uint ne12; uint ne13; + float sf0; float sf1; float sf2; float sf3; +} p; + +#include "types.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (idx >= p.ne) { + return; + } + + const uint i10 = idx % p.ne10; + const uint i11 = (idx / p.ne10) % p.ne11; + const uint i12 = (idx / (p.ne10 * p.ne11)) % p.ne12; + const uint i13 = (idx / (p.ne10 * p.ne11 * p.ne12)) % p.ne13; + + const uint i00 = uint(i10 / p.sf0); + const uint i01 = uint(i11 / p.sf1); + const uint i02 = uint(i12 / p.sf2); + const uint i03 = uint(i13 / p.sf3); + + data_d[p.d_offset + idx] = D_TYPE(data_a[p.a_offset + i03 * p.nb03 + i02 * p.nb02 + i01 * p.nb01 + i00 * p.nb00]); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp new file mode 100644 index 000000000..8111c0638 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -0,0 +1,594 @@ + + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef _WIN32 + #include + #include // For _mkdir on Windows + #include // For std::replace on w64devkit +#else + #include + #include + #include +#endif + +#include + +#define ASYNCIO_CONCURRENCY 64 + +std::mutex lock; +std::vector> shader_fnames; + +std::string GLSLC = "glslc"; +std::string input_dir = "vulkan-shaders"; +std::string output_dir = "/tmp"; +std::string target_hpp = "ggml-vulkan-shaders.hpp"; +std::string target_cpp = "ggml-vulkan-shaders.cpp"; +bool no_clean = false; + +const std::vector type_names = { + "f32", + "f16", + "q4_0", + "q4_1", + "q5_0", + "q5_1", + "q8_0", + "q2_k", + "q3_k", + "q4_k", + "q5_k", + "q6_k", + "iq4_nl" +}; + +namespace { +void execute_command(const std::string& command, std::string& stdout_str, std::string& stderr_str) { +#ifdef _WIN32 + HANDLE stdout_read, stdout_write; + HANDLE stderr_read, stderr_write; + SECURITY_ATTRIBUTES sa = { sizeof(SECURITY_ATTRIBUTES), NULL, TRUE }; + + if (!CreatePipe(&stdout_read, &stdout_write, &sa, 0) || + !SetHandleInformation(stdout_read, HANDLE_FLAG_INHERIT, 0)) { + throw std::runtime_error("Failed to create stdout pipe"); + } + + if (!CreatePipe(&stderr_read, &stderr_write, &sa, 0) || + !SetHandleInformation(stderr_read, HANDLE_FLAG_INHERIT, 0)) { + throw std::runtime_error("Failed to create stderr pipe"); + } + + PROCESS_INFORMATION pi; + STARTUPINFOA si = {}; + si.cb = sizeof(STARTUPINFOA); + si.dwFlags = STARTF_USESTDHANDLES; + si.hStdOutput = stdout_write; + si.hStdError = stderr_write; + + std::vector cmd(command.begin(), command.end()); + cmd.push_back('\0'); + + if (!CreateProcessA(NULL, cmd.data(), NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi)) { + throw std::runtime_error("Failed to create process"); + } + + CloseHandle(stdout_write); + CloseHandle(stderr_write); + + std::array buffer; + DWORD bytes_read; + + while (ReadFile(stdout_read, buffer.data(), (DWORD)buffer.size(), &bytes_read, NULL) && bytes_read > 0) { + stdout_str.append(buffer.data(), bytes_read); + } + + while (ReadFile(stderr_read, buffer.data(), (DWORD)buffer.size(), &bytes_read, NULL) && bytes_read > 0) { + stderr_str.append(buffer.data(), bytes_read); + } + + CloseHandle(stdout_read); + CloseHandle(stderr_read); + WaitForSingleObject(pi.hProcess, INFINITE); + CloseHandle(pi.hProcess); + CloseHandle(pi.hThread); +#else +int stdout_pipe[2]; + int stderr_pipe[2]; + + if (pipe(stdout_pipe) != 0 || pipe(stderr_pipe) != 0) { + throw std::runtime_error("Failed to create pipes"); + } + + pid_t pid = fork(); + if (pid < 0) { + throw std::runtime_error("Failed to fork process"); + } + + if (pid == 0) { + close(stdout_pipe[0]); + close(stderr_pipe[0]); + dup2(stdout_pipe[1], STDOUT_FILENO); + dup2(stderr_pipe[1], STDERR_FILENO); + close(stdout_pipe[1]); + close(stderr_pipe[1]); + execl("/bin/sh", "sh", "-c", command.c_str(), (char*) nullptr); + _exit(EXIT_FAILURE); + } else { + close(stdout_pipe[1]); + close(stderr_pipe[1]); + + std::array buffer; + ssize_t bytes_read; + + while ((bytes_read = read(stdout_pipe[0], buffer.data(), buffer.size())) > 0) { + stdout_str.append(buffer.data(), bytes_read); + } + + while ((bytes_read = read(stderr_pipe[0], buffer.data(), buffer.size())) > 0) { + stderr_str.append(buffer.data(), bytes_read); + } + + close(stdout_pipe[0]); + close(stderr_pipe[0]); + waitpid(pid, nullptr, 0); + } +#endif +} + +bool directory_exists(const std::string& path) { + struct stat info; + if (stat(path.c_str(), &info) != 0) { + return false; // Path doesn't exist or can't be accessed + } + return (info.st_mode & S_IFDIR) != 0; // Check if it is a directory +} + +bool create_directory(const std::string& path) { +#ifdef _WIN32 + return _mkdir(path.c_str()) == 0 || errno == EEXIST; // EEXIST means the directory already exists +#else + return mkdir(path.c_str(), 0755) == 0 || errno == EEXIST; // 0755 is the directory permissions +#endif +} + +std::string to_uppercase(const std::string& input) { + std::string result = input; + for (char& c : result) { + c = std::toupper(c); + } + return result; +} + +bool string_ends_with(const std::string& str, const std::string& suffix) { + if (suffix.size() > str.size()) { + return false; + } + return std::equal(suffix.rbegin(), suffix.rend(), str.rbegin()); +} + +static const char path_separator = '/'; + +std::string join_paths(const std::string& path1, const std::string& path2) { + return path1 + path_separator + path2; +} + +std::string basename(const std::string &path) { + return path.substr(path.find_last_of("/\\") + 1); +} + +// variables to track number of compiles in progress +static uint32_t compile_count = 0; +static std::mutex compile_count_mutex; +static std::condition_variable compile_count_cond; + +void string_to_spv_func(const std::string& _name, const std::string& in_fname, const std::map& defines, bool fp16 = true, bool coopmat = false, bool coopmat2 = false, bool f16acc = false) { + std::string name = _name + (f16acc ? "_f16acc" : "") + (coopmat ? "_coopmat" : "") + (coopmat2 ? "_cm2" : (fp16 ? "" : "_fp32")); + std::string out_fname = join_paths(output_dir, name + ".spv"); + std::string in_path = join_paths(input_dir, in_fname); + + std::string target_env = (name.find("_cm2") != std::string::npos) ? "--target-env=vulkan1.3" : "--target-env=vulkan1.2"; + + // disable spirv-opt for coopmat shaders for https://github.com/ggerganov/llama.cpp/issues/10734 + std::string opt_level = coopmat ? "" : "-O"; + + #ifdef _WIN32 + std::vector cmd = {GLSLC, "-fshader-stage=compute", target_env, opt_level, "\"" + in_path + "\"", "-o", "\"" + out_fname + "\""}; + #else + std::vector cmd = {GLSLC, "-fshader-stage=compute", target_env, opt_level, in_path, "-o", out_fname}; + #endif + + #ifdef GGML_VULKAN_SHADER_DEBUG_INFO + cmd.push_back("-g"); + #endif + + for (const auto& define : defines) { + cmd.push_back("-D" + define.first + "=" + define.second); + } + + std::string command; + for (const auto& part : cmd) { + command += part + " "; + } + + std::string stdout_str, stderr_str; + try { + // std::cout << "Executing command: "; + // for (const auto& part : cmd) { + // std::cout << part << " "; + // } + // std::cout << std::endl; + + execute_command(command, stdout_str, stderr_str); + if (!stderr_str.empty()) { + std::cerr << "cannot compile " << name << "\n\n" << command << "\n\n" << stderr_str << std::endl; + return; + } + + std::lock_guard guard(lock); + shader_fnames.push_back(std::make_pair(name, out_fname)); + } catch (const std::exception& e) { + std::cerr << "Error executing command for " << name << ": " << e.what() << std::endl; + } + { + std::lock_guard guard(compile_count_mutex); + assert(compile_count > 0); + compile_count--; + } + compile_count_cond.notify_all(); +} + +std::map merge_maps(const std::map& a, const std::map& b) { + std::map result = a; + result.insert(b.begin(), b.end()); + return result; +} + +static std::vector> compiles; +void string_to_spv(const std::string& _name, const std::string& in_fname, const std::map& defines, bool fp16 = true, bool coopmat = false, bool coopmat2 = false, bool f16acc = false) { + { + // wait until fewer than N compiles are in progress. + // 16 is an arbitrary limit, the goal is to avoid "failed to create pipe" errors. + uint32_t N = 16; + std::unique_lock guard(compile_count_mutex); + while (compile_count >= N) { + compile_count_cond.wait(guard); + } + compile_count++; + } + compiles.push_back(std::async(string_to_spv_func, _name, in_fname, defines, fp16, coopmat, coopmat2, f16acc)); +} + +void matmul_shaders(bool fp16, bool matmul_id, bool coopmat, bool coopmat2, bool f16acc) { + std::string load_vec = coopmat2 ? "1" : fp16 ? "8" : "4"; + std::string aligned_b_type_f32 = coopmat2 ? "float" : fp16 ? "mat2x4" : "vec4"; + std::string aligned_b_type_f16 = coopmat2 ? "float16_t" : fp16 ? "f16mat2x4" : "f16vec4"; + + std::map base_dict = {{"FLOAT_TYPE", (coopmat2 || fp16) ? "float16_t" : "float"}}; + std::string shader_name = "matmul"; + + if (matmul_id) { + base_dict["MUL_MAT_ID"] = "1"; + shader_name = "matmul_id"; + } + + if (fp16) { + base_dict["FLOAT16"] = "1"; + } + + base_dict["ACC_TYPE"] = f16acc ? "float16_t" : "float"; + + if (coopmat) { + base_dict["COOPMAT"] = "1"; + } + + base_dict["ACC_TYPE"] = f16acc ? "float16_t" : "float"; + + std::string source_name = coopmat2 ? "mul_mm_cm2.comp" : "mul_mm.comp"; + + // Shaders with f16 B_TYPE + string_to_spv(shader_name + "_f32_f16", source_name, merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}, }), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_f32_f16_aligned", source_name, merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + + string_to_spv(shader_name + "_f16_aligned", source_name, merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_f16", source_name, merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); + + for (const auto& tname : type_names) { + std::string data_a_key = "DATA_A_" + to_uppercase(tname); + // For unaligned, load one at a time for f32/f16, or two at a time for quants + std::string load_vec_a_unaligned = (coopmat2 || tname == "f32" || tname == "f16") ? "1" : "2"; + // For aligned matmul loads + std::string load_vec_a = (coopmat2 || tname == "f32" || tname == "f16") ? load_vec : "2"; + + string_to_spv(shader_name + "_" + tname + "_f32", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_" + tname + "_f32_aligned", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + + if (tname != "f16" && tname != "f32") { + string_to_spv(shader_name + "_" + tname + "_f16", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_" + tname + "_f16_aligned", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + } + } +} + +void process_shaders() { + std::cout << "ggml_vulkan: Generating and compiling shaders to SPIR-V" << std::endl; + std::map base_dict = {{"FLOAT_TYPE", "float"}}; + + // matmul + for (const auto& matmul_id : {false, true}) { + // No coopmats + // fp32 + matmul_shaders(false, matmul_id, false, false, false); + + // fp16, fp32acc and fp16acc + matmul_shaders(true, matmul_id, false, false, false); + matmul_shaders(true, matmul_id, false, false, true); + + // Coopmat, fp32acc and fp16acc + matmul_shaders(true, matmul_id, true, false, false); + matmul_shaders(true, matmul_id, true, false, true); + +#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + // Coopmat2, fp32acc and fp16acc + matmul_shaders(true, matmul_id, false, true, false); + matmul_shaders(true, matmul_id, false, true, true); +#endif + } + +#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + // flash attention + for (const auto& f16acc : {false, true}) { + std::string acctype = f16acc ? "float16_t" : "float"; + + for (const auto& tname : type_names) { + if (tname == "f32") { + continue; + } + + if (tname == "f16") { + string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm2.comp", + merge_maps(base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"ACC_TYPE", acctype}}), true, false, true, f16acc); + } else { + std::string data_a_key = "DATA_A_" + to_uppercase(tname); + string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm2.comp", + merge_maps(base_dict, {{data_a_key, "1"}, {"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"ACC_TYPE", acctype}, {"DEQUANTFUNC", "dequantFunc"+to_uppercase(tname) }, {"BLOCK_SIZE", "QUANT_K_"+to_uppercase(tname) }}), true, false, true, f16acc); + } + } + } +#endif + + for (const auto& tname : type_names) { + // mul mat vec + std::string data_a_key = "DATA_A_" + to_uppercase(tname); + std::string shader = (string_ends_with(tname, "_k")) ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp"; + + string_to_spv("mul_mat_vec_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPE_VEC2", "vec2"}, {"B_TYPE_VEC4", "vec4"}, {"D_TYPE", "float"}})); + string_to_spv("mul_mat_vec_" + tname + "_f16_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"B_TYPE_VEC2", "f16vec2"}, {"B_TYPE_VEC4", "f16vec4"}, {"D_TYPE", "float"}})); + + string_to_spv("mul_mat_vec_id_" + tname + "_f32", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPE_VEC2", "vec2"}, {"B_TYPE_VEC4", "vec4"}, {"D_TYPE", "float"}})); + + // Dequant shaders + if (tname != "f16") { + string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}})); + } + + if (!string_ends_with(tname, "_k")) { + shader = (tname == "f32" || tname == "f16") ? "get_rows.comp" : "get_rows_quant.comp"; + + if (tname == "f16") { + string_to_spv("get_rows_" + tname, shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}})); + } else { + string_to_spv("get_rows_" + tname, shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}})); + } + string_to_spv("get_rows_" + tname + "_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float"}})); + } + } + + string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("mul_mat_vec_nc_f16_f32", "mul_mat_vec_nc.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); + + // Norms + string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + + string_to_spv("cpy_f32_f32", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("cpy_f32_f16", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); + string_to_spv("cpy_f16_f16", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); + string_to_spv("contig_cpy_f32_f32", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("contig_cpy_f32_f16", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); + string_to_spv("contig_cpy_f16_f16", "contig_copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); + + string_to_spv("add_f32", "add.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + string_to_spv("add_f16_f32_f16", "add.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("acc_f32", "acc.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("split_k_reduce", "mul_mat_split_k_reduce.comp", {}); + + string_to_spv("mul_f32", "mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("div_f32", "div.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("repeat_f32", "repeat.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + + string_to_spv("scale_f32", "scale.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("sqr_f32", "square.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("sin_f32", "sin.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("cos_f32", "cos.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("clamp_f32", "clamp.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("pad_f32", "pad.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + + string_to_spv("concat_f32", "concat.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("concat_f16", "concat.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); + string_to_spv("concat_i32", "concat.comp", {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}}); + + string_to_spv("upscale_f32", "upscale.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); + + string_to_spv("gelu_f32", "gelu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("gelu_quick_f32", "gelu_quick.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("silu_f32", "silu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("relu_f32", "relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("leaky_relu_f32", "leaky_relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("tanh_f32", "tanh.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + + string_to_spv("diag_mask_inf_f32", "diag_mask_inf.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + + string_to_spv("soft_max_f32", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("soft_max_f32_f16", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}})); + + string_to_spv("rope_norm_f32", "rope_norm.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("rope_norm_f16", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); + string_to_spv("rope_norm_f16_rte", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}}); + + string_to_spv("rope_neox_f32", "rope_neox.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("rope_neox_f16", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); + string_to_spv("rope_neox_f16_rte", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}}); + + string_to_spv("argsort_f32", "argsort.comp", {{"A_TYPE", "float"}}); + + string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + + string_to_spv("im2col_f32", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("im2col_f32_f16", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}})); + string_to_spv("im2col_f32_f16_rte", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}})); + + string_to_spv("timestep_embedding_f32", "timestep_embedding.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + + string_to_spv("pool2d_f32", "pool2d.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + + string_to_spv("rwkv_wkv6_f32", "wkv6.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); + + for (auto &c : compiles) { + c.wait(); + } +} + +void write_output_files() { + FILE* hdr = fopen(target_hpp.c_str(), "w"); + FILE* src = fopen(target_cpp.c_str(), "w"); + + fprintf(hdr, "#include \n\n"); + fprintf(src, "#include \"%s\"\n\n", basename(target_hpp).c_str()); + + for (const auto& pair : shader_fnames) { + const std::string& name = pair.first; + #ifdef _WIN32 + std::string path = pair.second; + std::replace(path.begin(), path.end(), '/', '\\' ); + #else + const std::string& path = pair.second; + #endif + + FILE* spv = fopen(path.c_str(), "rb"); + if (!spv) { + std::cerr << "Error opening SPIR-V file: " << path << " (" << strerror(errno) << ")\n"; + continue; + } + + fseek(spv, 0, SEEK_END); + size_t size = ftell(spv); + fseek(spv, 0, SEEK_SET); + + std::vector data(size); + size_t read_size = fread(data.data(), 1, size, spv); + fclose(spv); + if (read_size != size) { + std::cerr << "Error reading SPIR-V file: " << path << " (" << strerror(errno) << ")\n"; + continue; + } + + fprintf(hdr, "extern unsigned char %s_data[%zu];\n", name.c_str(), size); + fprintf(hdr, "const uint64_t %s_len = %zu;\n\n", name.c_str(), size); + + fprintf(src, "unsigned char %s_data[%zu] = {\n", name.c_str(), size); + for (size_t i = 0; i < size; ++i) { + fprintf(src, "0x%02x,", data[i]); + if ((i + 1) % 12 == 0) fprintf(src, "\n"); + } + fprintf(src, "\n};\n\n"); + + if (!no_clean) { + std::remove(path.c_str()); + } + } + + fclose(hdr); + fclose(src); +} +} + +int main(int argc, char** argv) { + std::map args; + for (int i = 1; i < argc; ++i) { + std::string arg = argv[i]; + if (arg.rfind("--", 0) == 0) { + if (i + 1 < argc && argv[i + 1][0] != '-') { + args[arg] = argv[i + 1]; + ++i; + } else { + args[arg] = ""; + } + } + } + + if (args.find("--glslc") != args.end()) { + GLSLC = args["--glslc"]; // Path to glslc + } + if (args.find("--input-dir") != args.end()) { + input_dir = args["--input-dir"]; // Directory containing shader sources + } + if (args.find("--output-dir") != args.end()) { + output_dir = args["--output-dir"]; // Directory for containing SPIR-V output + } + if (args.find("--target-hpp") != args.end()) { + target_hpp = args["--target-hpp"]; // Path to generated header file + } + if (args.find("--target-cpp") != args.end()) { + target_cpp = args["--target-cpp"]; // Path to generated cpp file + } + if (args.find("--no-clean") != args.end()) { + no_clean = true; // Keep temporary SPIR-V files in output-dir after build + } + + if (!directory_exists(input_dir)) { + std::cerr << "\"" << input_dir << "\" must be a valid directory containing shader sources" << std::endl; + return EXIT_FAILURE; + } + + if (!directory_exists(output_dir)) { + if (!create_directory(output_dir)) { + std::cerr << "Error creating output directory: " << output_dir << "\n"; + return EXIT_FAILURE; + } + } + + process_shaders(); + + write_output_files(); + + return EXIT_SUCCESS; +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp new file mode 100644 index 000000000..35cc6c45f --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp @@ -0,0 +1,87 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +#define BLOCK_SIZE 64 +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform Parameters { + uint B; + uint T; + uint C; + uint H; +}; + +layout(binding = 0) readonly buffer KBuf { A_TYPE k[]; }; +layout(binding = 1) readonly buffer VBuf { A_TYPE v[]; }; +layout(binding = 2) readonly buffer RBuf { A_TYPE r[]; }; +layout(binding = 3) readonly buffer TimeFBuf { A_TYPE tf[]; }; +layout(binding = 4) readonly buffer TimeDBuf { A_TYPE td[]; }; +layout(binding = 5) readonly buffer StateBuf { A_TYPE state_in[]; }; +layout(binding = 6) buffer DstBuf { A_TYPE dst[]; }; + +shared A_TYPE _k[BLOCK_SIZE], _r[BLOCK_SIZE], _tf[BLOCK_SIZE], _td[BLOCK_SIZE]; + +void main() { + const uint head_size = BLOCK_SIZE; + const uint batch_id = gl_WorkGroupID.x / H; + const uint head_id = gl_WorkGroupID.x % H; + const uint tid = gl_LocalInvocationID.x; + + const uint state_size = C * head_size; + const uint n_seq_tokens = T / B; + + if (batch_id >= B || head_id >= H) { + return; + } + + A_TYPE state[BLOCK_SIZE]; + [[unroll]] for (uint i = 0; i < head_size; i++) { + state[i] = state_in[batch_id * state_size + head_id * head_size * head_size + + i * head_size + tid]; + } + + barrier(); + _tf[tid] = tf[head_id * head_size + tid]; + barrier(); + + const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid; + const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid; + + for (uint t = start_t; t < end_t; t += C) { + barrier(); + _k[tid] = k[t]; + _r[tid] = r[t]; + _td[tid] = td[t]; + barrier(); + + const A_TYPE v_val = v[t]; + A_TYPE y = 0.0; + + [[unroll]] for (uint j = 0; j < head_size; j += 4) { + vec4 k_vec = vec4(_k[j], _k[j+1], _k[j+2], _k[j+3]); + vec4 r_vec = vec4(_r[j], _r[j+1], _r[j+2], _r[j+3]); + vec4 tf_vec = vec4(_tf[j], _tf[j+1], _tf[j+2], _tf[j+3]); + vec4 td_vec = vec4(_td[j], _td[j+1], _td[j+2], _td[j+3]); + vec4 s_vec = vec4(state[j], state[j+1], state[j+2], state[j+3]); + + vec4 kv = k_vec * v_val; + + vec4 temp = tf_vec * kv + s_vec; + y += dot(r_vec, temp); + + s_vec = s_vec * td_vec + kv; + state[j] = s_vec.x; + state[j+1] = s_vec.y; + state[j+2] = s_vec.z; + state[j+3] = s_vec.w; + } + + dst[t] = y; + } + + [[unroll]] for (uint i = 0; i < head_size; i++) { + dst[T * C + batch_id * state_size + head_id * head_size * head_size + + i * head_size + tid] = state[i]; + } +} From 189cbb40a6f7f3ea55d1986827f5912a10f485ea Mon Sep 17 00:00:00 2001 From: Vadim Grinco Date: Sat, 8 Mar 2025 19:40:53 +0100 Subject: [PATCH 22/29] Updated dockerfile https://github.com/whyvl/ollama-vulkan/issues/7#issuecomment-2660836871 Signed-off-by: Vadim Grinco --- Dockerfile | 181 ++++++++++++++++++++--------------------------------- 1 file changed, 69 insertions(+), 112 deletions(-) diff --git a/Dockerfile b/Dockerfile index 4136fca71..4e87ba43c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,131 +1,88 @@ -# vim: filetype=dockerfile +FROM --platform=linux/amd64 library/ubuntu:noble as builder -ARG FLAVOR=${TARGETARCH} +ENV DEBIAN_FRONTEND="noninteractive" -ARG ROCMVERSION=6.3.3 -ARG JETPACK5VERSION=r35.4.1 -ARG JETPACK6VERSION=r36.4.0 -ARG CMAKEVERSION=3.31.2 +ENV VULKAN_VER_BASE="1.3.296" +ENV VULKAN_VER="${VULKAN_VER_BASE}.0" +ENV UBUNTU_VERSION="noble" -# CUDA v11 requires gcc v10. v10.3 has regressions, so the rockylinux 8.5 AppStream has the latest compatible version -FROM --platform=linux/amd64 rocm/dev-almalinux-8:${ROCMVERSION}-complete AS base-amd64 -RUN yum install -y yum-utils \ - && yum-config-manager --add-repo https://dl.rockylinux.org/vault/rocky/8.5/AppStream/\$basearch/os/ \ - && rpm --import https://dl.rockylinux.org/pub/rocky/RPM-GPG-KEY-Rocky-8 \ - && dnf install -y yum-utils ccache gcc-toolset-10-gcc-10.2.1-8.2.el8 gcc-toolset-10-gcc-c++-10.2.1-8.2.el8 gcc-toolset-10-binutils-2.35-11.el8 \ - && yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo -ENV PATH=/opt/rh/gcc-toolset-10/root/usr/bin:$PATH - -FROM --platform=linux/arm64 almalinux:8 AS base-arm64 -# install epel-release for ccache -RUN yum install -y yum-utils epel-release \ - && dnf install -y clang ccache \ - && yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/sbsa/cuda-rhel8.repo -ENV CC=clang CXX=clang++ - -FROM base-${TARGETARCH} AS base -ARG CMAKEVERSION -RUN curl -fsSL https://github.com/Kitware/CMake/releases/download/v${CMAKEVERSION}/cmake-${CMAKEVERSION}-linux-$(uname -m).tar.gz | tar xz -C /usr/local --strip-components 1 -COPY CMakeLists.txt CMakePresets.json . -COPY ml/backend/ggml/ggml ml/backend/ggml/ggml +ENV GOLANG_VERSION="1.22.8" +ENV GOARCH="amd64" +ENV CGO_ENABLED=1 ENV LDFLAGS=-s -FROM base AS cpu -RUN dnf install -y gcc-toolset-11-gcc gcc-toolset-11-gcc-c++ -ENV PATH=/opt/rh/gcc-toolset-11/root/usr/bin:$PATH -RUN --mount=type=cache,target=/root/.ccache \ - cmake --preset 'CPU' \ - && cmake --build --parallel --preset 'CPU' \ - && cmake --install build --component CPU --strip --parallel 8 +# Default mirror was very slow +RUN \ + sed -i 's/archive.ubuntu.com/gb.archive.ubuntu.com/g' /etc/apt/sources.list.d/ubuntu.sources -FROM base AS cuda-11 -ARG CUDA11VERSION=11.3 -RUN dnf install -y cuda-toolkit-${CUDA11VERSION//./-} -ENV PATH=/usr/local/cuda-11/bin:$PATH -RUN --mount=type=cache,target=/root/.ccache \ - cmake --preset 'CUDA 11' \ - && cmake --build --parallel --preset 'CUDA 11' \ - && cmake --install build --component CUDA --strip --parallel 8 +RUN \ + apt-get update && \ + apt-get install -y ca-certificates build-essential ccache cmake wget git curl rsync xz-utils libcap-dev -FROM base AS cuda-12 -ARG CUDA12VERSION=12.8 -RUN dnf install -y cuda-toolkit-${CUDA12VERSION//./-} -ENV PATH=/usr/local/cuda-12/bin:$PATH -RUN --mount=type=cache,target=/root/.ccache \ - cmake --preset 'CUDA 12' \ - && cmake --build --parallel --preset 'CUDA 12' \ - && cmake --install build --component CUDA --strip --parallel 8 +RUN \ + mkdir -p /usr/local 2>/dev/null || true && \ + curl -s -L https://dl.google.com/go/go${GOLANG_VERSION}.linux-${GOARCH}.tar.gz | tar -xz -C /usr/local && \ + ln -s /usr/local/go/bin/go /usr/local/bin/go && \ + ln -s /usr/local/go/bin/gofmt /usr/local/bin/gofmt -FROM base AS rocm-6 -ENV PATH=/opt/rocm/hcc/bin:/opt/rocm/hip/bin:/opt/rocm/bin:/opt/rocm/hcc/bin:$PATH -RUN --mount=type=cache,target=/root/.ccache \ - cmake --preset 'ROCm 6' \ - && cmake --build --parallel --preset 'ROCm 6' \ - && cmake --install build --component HIP --strip --parallel 8 -FROM --platform=linux/arm64 nvcr.io/nvidia/l4t-jetpack:${JETPACK5VERSION} AS jetpack-5 -ARG CMAKEVERSION -RUN apt-get update && apt-get install -y curl ccache \ - && curl -fsSL https://github.com/Kitware/CMake/releases/download/v${CMAKEVERSION}/cmake-${CMAKEVERSION}-linux-$(uname -m).tar.gz | tar xz -C /usr/local --strip-components 1 -COPY CMakeLists.txt CMakePresets.json . -COPY ml/backend/ggml/ggml ml/backend/ggml/ggml -RUN --mount=type=cache,target=/root/.ccache \ - cmake --preset 'JetPack 5' \ - && cmake --build --parallel --preset 'JetPack 5' \ - && cmake --install build --component CUDA --strip --parallel 8 +RUN \ + wget -qO - https://packages.lunarg.com/lunarg-signing-key-pub.asc | gpg --dearmor -o /etc/apt/trusted.gpg.d/lunarg-signing-key-pub.gpg && \ + wget -qO /etc/apt/sources.list.d/lunarg-vulkan-${UBUNTU_VERSION}.list https://packages.lunarg.com/vulkan/${VULKAN_VER_BASE}/lunarg-vulkan-${VULKAN_VER_BASE}-${UBUNTU_VERSION}.list && \ + apt update && apt install -y vulkan-sdk -FROM --platform=linux/arm64 nvcr.io/nvidia/l4t-jetpack:${JETPACK6VERSION} AS jetpack-6 -ARG CMAKEVERSION -RUN apt-get update && apt-get install -y curl ccache \ - && curl -fsSL https://github.com/Kitware/CMake/releases/download/v${CMAKEVERSION}/cmake-${CMAKEVERSION}-linux-$(uname -m).tar.gz | tar xz -C /usr/local --strip-components 1 -COPY CMakeLists.txt CMakePresets.json . -COPY ml/backend/ggml/ggml ml/backend/ggml/ggml -RUN --mount=type=cache,target=/root/.ccache \ - cmake --preset 'JetPack 6' \ - && cmake --build --parallel --preset 'JetPack 6' \ - && cmake --install build --component CUDA --strip --parallel 8 +# Last testet ollama-vulkan commit: +# 2d443b3dd660a1fd2760d64538512df93648b4bb +COPY patches/ /tmp/patches/ +RUN \ + git clone https://github.com/pufferffish/ollama-vulkan.git "/tmp/ollama-vulkan-git" && \ + cd "/tmp/ollama-vulkan-git" && \ + git checkout 2d443b3dd660a1fd2760d64538512df93648b4bb && git checkout -b ollama_vulkan_stable && \ + git config user.name "Builder" && git config user.email "builder@local" && \ + git remote add ollama_vanilla https://github.com/ollama/ollama.git && \ + git fetch ollama_vanilla --tags && git checkout v0.5.11 && git checkout -b ollama_vanilla_stable && \ + git checkout ollama_vulkan_stable && git merge ollama_vanilla_stable --allow-unrelated-histories --no-edit && \ + for p in /tmp/patches/00-fix-vulkan-building.patch; do patch -p1 < $p; done -FROM base AS build -WORKDIR /go/src/github.com/ollama/ollama -COPY go.mod go.sum . -RUN curl -fsSL https://golang.org/dl/go$(awk '/^go/ { print $2 }' go.mod).linux-$(case $(uname -m) in x86_64) echo amd64 ;; aarch64) echo arm64 ;; esac).tar.gz | tar xz -C /usr/local -ENV PATH=/usr/local/go/bin:$PATH -RUN go mod download -COPY . . -ARG GOFLAGS="'-ldflags=-w -s'" -ENV CGO_ENABLED=1 -RUN --mount=type=cache,target=/root/.cache/go-build \ - go build -trimpath -buildmode=pie -o /bin/ollama . +RUN \ + cd "/tmp/ollama-vulkan-git" && \ + make -f Makefile.sync clean sync -FROM --platform=linux/amd64 scratch AS amd64 -COPY --from=cuda-11 dist/lib/ollama/cuda_v11 /lib/ollama/cuda_v11 -COPY --from=cuda-12 dist/lib/ollama/cuda_v12 /lib/ollama/cuda_v12 -FROM --platform=linux/arm64 scratch AS arm64 -COPY --from=cuda-11 dist/lib/ollama/cuda_v11 /lib/ollama/cuda_v11 -COPY --from=cuda-12 dist/lib/ollama/cuda_v12 /lib/ollama/cuda_v12 -COPY --from=jetpack-5 dist/lib/ollama/cuda_v11 lib/ollama/cuda_jetpack5 -COPY --from=jetpack-6 dist/lib/ollama/cuda_v12 lib/ollama/cuda_jetpack6 +FROM builder AS cpu-build +RUN \ + cd "/tmp/ollama-vulkan-git" && \ + cmake --preset CPU && cmake --build --parallel --preset CPU && \ + cmake --install build --component CPU --strip -FROM scratch AS rocm -COPY --from=rocm-6 dist/lib/ollama/rocm /lib/ollama/rocm +FROM builder AS vulkan-build +RUN \ + cd "/tmp/ollama-vulkan-git" && \ + cmake --preset Vulkan && \ + cmake --build --parallel --preset Vulkan && \ + cmake --install build --component Vulkan --strip -FROM ${FLAVOR} AS archive -COPY --from=cpu dist/lib/ollama /lib/ollama -COPY --from=build /bin/ollama /bin/ollama +FROM builder AS binary-build +RUN \ + cd "/tmp/ollama-vulkan-git" && \ + . scripts/env.sh && \ + mkdir -p dist/bin && \ + go build -trimpath -buildmode=pie -o dist/bin/ollama . + + +FROM --platform=linux/amd64 library/ubuntu:noble +RUN \ + apt-get update && \ + apt-get install -y ca-certificates libcap2 libvulkan1 && \ + apt-get clean && rm -rf /var/lib/apt/lists/* +COPY --from=cpu-build /tmp/ollama-vulkan-git/dist/lib/ollama/ /lib/ollama/ +COPY --from=vulkan-build /tmp/ollama-vulkan-git/dist/lib/ollama/vulkan/ /lib/ollama/vulkan/ +COPY --from=binary-build /tmp/ollama-vulkan-git/dist/bin/ /bin/ + +RUN find /lib/ollama && find /bin/ollama -FROM ubuntu:20.04 -RUN apt-get update \ - && apt-get install -y ca-certificates \ - && apt-get clean \ - && rm -rf /var/lib/apt/lists/* -COPY --from=archive /bin /usr/bin -ENV PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin -COPY --from=archive /lib/ollama /usr/lib/ollama -ENV LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64 -ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility -ENV NVIDIA_VISIBLE_DEVICES=all -ENV OLLAMA_HOST=0.0.0.0:11434 EXPOSE 11434 +ENV OLLAMA_HOST 0.0.0.0 + ENTRYPOINT ["/bin/ollama"] CMD ["serve"] From 81465ca37494217f0d6f074eb47306835d26638b Mon Sep 17 00:00:00 2001 From: Vadim Grinco Date: Sun, 9 Mar 2025 20:42:32 +0100 Subject: [PATCH 23/29] Installing rocm library Signed-off-by: Vadim Grinco --- Dockerfile | 105 ++++++++++++++++++++++++++++------------------------- 1 file changed, 55 insertions(+), 50 deletions(-) diff --git a/Dockerfile b/Dockerfile index 4e87ba43c..25d8ddd3e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,88 +1,93 @@ -FROM --platform=linux/amd64 library/ubuntu:noble as builder +# Base Image +FROM --platform=linux/amd64 library/ubuntu:noble AS builder +# Set Environment Variables ENV DEBIAN_FRONTEND="noninteractive" - ENV VULKAN_VER_BASE="1.3.296" ENV VULKAN_VER="${VULKAN_VER_BASE}.0" ENV UBUNTU_VERSION="noble" - ENV GOLANG_VERSION="1.22.8" ENV GOARCH="amd64" ENV CGO_ENABLED=1 ENV LDFLAGS=-s -# Default mirror was very slow -RUN \ - sed -i 's/archive.ubuntu.com/gb.archive.ubuntu.com/g' /etc/apt/sources.list.d/ubuntu.sources +# Set up faster package mirrors +RUN sed -i 's/archive.ubuntu.com/gb.archive.ubuntu.com/g' /etc/apt/sources.list.d/ubuntu.sources -RUN \ - apt-get update && \ - apt-get install -y ca-certificates build-essential ccache cmake wget git curl rsync xz-utils libcap-dev +# Install Required Dependencies +RUN apt-get update && apt-get install -y \ + ca-certificates build-essential ccache cmake wget git curl rsync xz-utils libcap-dev \ + && apt-get clean && rm -rf /var/lib/apt/lists/* -RUN \ - mkdir -p /usr/local 2>/dev/null || true && \ +# Install Go +RUN mkdir -p /usr/local && \ curl -s -L https://dl.google.com/go/go${GOLANG_VERSION}.linux-${GOARCH}.tar.gz | tar -xz -C /usr/local && \ ln -s /usr/local/go/bin/go /usr/local/bin/go && \ ln -s /usr/local/go/bin/gofmt /usr/local/bin/gofmt - -RUN \ - wget -qO - https://packages.lunarg.com/lunarg-signing-key-pub.asc | gpg --dearmor -o /etc/apt/trusted.gpg.d/lunarg-signing-key-pub.gpg && \ +# Install Vulkan SDK +RUN wget -qO - https://packages.lunarg.com/lunarg-signing-key-pub.asc | gpg --dearmor -o /etc/apt/trusted.gpg.d/lunarg-signing-key-pub.gpg && \ wget -qO /etc/apt/sources.list.d/lunarg-vulkan-${UBUNTU_VERSION}.list https://packages.lunarg.com/vulkan/${VULKAN_VER_BASE}/lunarg-vulkan-${VULKAN_VER_BASE}-${UBUNTU_VERSION}.list && \ - apt update && apt install -y vulkan-sdk + apt update && apt install -y vulkan-sdk && \ + apt-get clean && rm -rf /var/lib/apt/lists/* -# Last testet ollama-vulkan commit: -# 2d443b3dd660a1fd2760d64538512df93648b4bb -COPY patches/ /tmp/patches/ -RUN \ - git clone https://github.com/pufferffish/ollama-vulkan.git "/tmp/ollama-vulkan-git" && \ - cd "/tmp/ollama-vulkan-git" && \ +# Install AMDVLK (Optional: If you want to use AMDVLK instead of RADV) +RUN wget -qO - http://repo.radeon.com/amdvlk/apt/debian/amdvlk.gpg.key | apt-key add && \ + echo "deb [arch=amd64,i386] http://repo.radeon.com/amdvlk/apt/debian/ bionic main" > /etc/apt/sources.list.d/amdvlk.list && \ + apt update && apt install -y amdvlk && \ + apt-get clean && rm -rf /var/lib/apt/lists/* + +# Set AMDVLK as the default Vulkan driver +ENV VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/amd_icd64.json + +# Clone Ollama Vulkan Fork +WORKDIR /opt +RUN git clone https://github.com/pufferffish/ollama-vulkan.git ollama-vulkan + +# Download and Apply Patches Automatically +WORKDIR /opt/ollama-vulkan +RUN mkdir -p patches && \ + wget -O patches/00-fix-vulkan-building.patch https://github.com/user-attachments/files/18783263/0002-fix-fix-vulkan-building.patch && \ git checkout 2d443b3dd660a1fd2760d64538512df93648b4bb && git checkout -b ollama_vulkan_stable && \ git config user.name "Builder" && git config user.email "builder@local" && \ git remote add ollama_vanilla https://github.com/ollama/ollama.git && \ - git fetch ollama_vanilla --tags && git checkout v0.5.11 && git checkout -b ollama_vanilla_stable && \ + git fetch ollama_vanilla --tags && git checkout v0.5.13 && git checkout -b ollama_vanilla_stable && \ git checkout ollama_vulkan_stable && git merge ollama_vanilla_stable --allow-unrelated-histories --no-edit && \ - for p in /tmp/patches/00-fix-vulkan-building.patch; do patch -p1 < $p; done + for p in patches/*.patch; do patch -p1 < $p; done -RUN \ - cd "/tmp/ollama-vulkan-git" && \ - make -f Makefile.sync clean sync +# Build Shared Libraries (CPU & Vulkan) +WORKDIR /opt/ollama-vulkan +RUN cmake -S . -B build -DCMAKE_BUILD_TYPE=Release +RUN cmake --build build --parallel +RUN cmake --install build --component CPU --strip +RUN cmake --install build --component Vulkan --strip +# Install rocm +RUN apt update +RUN apt install -y wget "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)" +RUN apt install -y python3-setuptools python3-wheel +RUN wget https://repo.radeon.com/amdgpu-install/6.3.3/ubuntu/noble/amdgpu-install_6.3.60303-1_all.deb -O /tmp/amdgpu-install_6.3.60303-1_all.deb +RUN apt install -y /tmp/amdgpu-install_6.3.60303-1_all.deb +RUN apt update && apt install -y rocm -FROM builder AS cpu-build -RUN \ - cd "/tmp/ollama-vulkan-git" && \ - cmake --preset CPU && cmake --build --parallel --preset CPU && \ - cmake --install build --component CPU --strip - -FROM builder AS vulkan-build -RUN \ - cd "/tmp/ollama-vulkan-git" && \ - cmake --preset Vulkan && \ - cmake --build --parallel --preset Vulkan && \ - cmake --install build --component Vulkan --strip - -FROM builder AS binary-build -RUN \ - cd "/tmp/ollama-vulkan-git" && \ +# Build Final Binary +RUN cd /opt/ollama-vulkan && \ . scripts/env.sh && \ mkdir -p dist/bin && \ go build -trimpath -buildmode=pie -o dist/bin/ollama . - +# Final Image FROM --platform=linux/amd64 library/ubuntu:noble -RUN \ - apt-get update && \ - apt-get install -y ca-certificates libcap2 libvulkan1 && \ +RUN apt-get update && apt-get install -y ca-certificates libcap2 libvulkan1 && \ apt-get clean && rm -rf /var/lib/apt/lists/* -COPY --from=cpu-build /tmp/ollama-vulkan-git/dist/lib/ollama/ /lib/ollama/ -COPY --from=vulkan-build /tmp/ollama-vulkan-git/dist/lib/ollama/vulkan/ /lib/ollama/vulkan/ -COPY --from=binary-build /tmp/ollama-vulkan-git/dist/bin/ /bin/ -RUN find /lib/ollama && find /bin/ollama +# Copy Built Components +COPY --from=builder /opt/ollama-vulkan/dist/bin/ollama /bin/ollama +# Expose Ollama Server Port EXPOSE 11434 ENV OLLAMA_HOST 0.0.0.0 +# Run Ollama Server ENTRYPOINT ["/bin/ollama"] CMD ["serve"] From 42bac5caddfb95c40c5d34d0fe123c4e3c0f1f2a Mon Sep 17 00:00:00 2001 From: Vadim Grinco Date: Sun, 9 Mar 2025 23:21:57 +0100 Subject: [PATCH 24/29] This version works well built based on this: https://github.com/whyvl/ollama-vulkan/issues/7#issuecomment-2660836871 Signed-off-by: Vadim Grinco --- Dockerfile | 116 +- patches/00-fix-vulkan-building.patch | 15297 +++++++++++++++++++++++++ 2 files changed, 15358 insertions(+), 55 deletions(-) create mode 100644 patches/00-fix-vulkan-building.patch diff --git a/Dockerfile b/Dockerfile index 25d8ddd3e..9e2928108 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,93 +1,99 @@ -# Base Image -FROM --platform=linux/amd64 library/ubuntu:noble AS builder +FROM --platform=linux/amd64 library/ubuntu:noble as builder -# Set Environment Variables ENV DEBIAN_FRONTEND="noninteractive" + ENV VULKAN_VER_BASE="1.3.296" ENV VULKAN_VER="${VULKAN_VER_BASE}.0" ENV UBUNTU_VERSION="noble" + ENV GOLANG_VERSION="1.22.8" ENV GOARCH="amd64" ENV CGO_ENABLED=1 ENV LDFLAGS=-s -# Set up faster package mirrors -RUN sed -i 's/archive.ubuntu.com/gb.archive.ubuntu.com/g' /etc/apt/sources.list.d/ubuntu.sources +# Default mirror was very slow +RUN \ + sed -i 's/archive.ubuntu.com/gb.archive.ubuntu.com/g' /etc/apt/sources.list.d/ubuntu.sources -# Install Required Dependencies -RUN apt-get update && apt-get install -y \ - ca-certificates build-essential ccache cmake wget git curl rsync xz-utils libcap-dev \ - && apt-get clean && rm -rf /var/lib/apt/lists/* +RUN \ + apt-get update && \ + apt-get install -y ca-certificates build-essential ccache cmake wget git curl rsync xz-utils libcap-dev -# Install Go -RUN mkdir -p /usr/local && \ +RUN \ + mkdir -p /usr/local 2>/dev/null || true && \ curl -s -L https://dl.google.com/go/go${GOLANG_VERSION}.linux-${GOARCH}.tar.gz | tar -xz -C /usr/local && \ ln -s /usr/local/go/bin/go /usr/local/bin/go && \ ln -s /usr/local/go/bin/gofmt /usr/local/bin/gofmt -# Install Vulkan SDK -RUN wget -qO - https://packages.lunarg.com/lunarg-signing-key-pub.asc | gpg --dearmor -o /etc/apt/trusted.gpg.d/lunarg-signing-key-pub.gpg && \ + +RUN \ + wget -qO - https://packages.lunarg.com/lunarg-signing-key-pub.asc | gpg --dearmor -o /etc/apt/trusted.gpg.d/lunarg-signing-key-pub.gpg && \ wget -qO /etc/apt/sources.list.d/lunarg-vulkan-${UBUNTU_VERSION}.list https://packages.lunarg.com/vulkan/${VULKAN_VER_BASE}/lunarg-vulkan-${VULKAN_VER_BASE}-${UBUNTU_VERSION}.list && \ - apt update && apt install -y vulkan-sdk && \ - apt-get clean && rm -rf /var/lib/apt/lists/* + apt update && apt install -y vulkan-sdk -# Install AMDVLK (Optional: If you want to use AMDVLK instead of RADV) -RUN wget -qO - http://repo.radeon.com/amdvlk/apt/debian/amdvlk.gpg.key | apt-key add && \ - echo "deb [arch=amd64,i386] http://repo.radeon.com/amdvlk/apt/debian/ bionic main" > /etc/apt/sources.list.d/amdvlk.list && \ - apt update && apt install -y amdvlk && \ - apt-get clean && rm -rf /var/lib/apt/lists/* - -# Set AMDVLK as the default Vulkan driver -ENV VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/amd_icd64.json - -# Clone Ollama Vulkan Fork -WORKDIR /opt -RUN git clone https://github.com/pufferffish/ollama-vulkan.git ollama-vulkan - -# Download and Apply Patches Automatically -WORKDIR /opt/ollama-vulkan -RUN mkdir -p patches && \ - wget -O patches/00-fix-vulkan-building.patch https://github.com/user-attachments/files/18783263/0002-fix-fix-vulkan-building.patch && \ +# Last testet ollama-vulkan commit: +# 2d443b3dd660a1fd2760d64538512df93648b4bb +COPY patches/ /tmp/patches/ +RUN \ + git clone https://github.com/pufferffish/ollama-vulkan.git "/tmp/ollama-vulkan-git" && \ + cd "/tmp/ollama-vulkan-git" && \ git checkout 2d443b3dd660a1fd2760d64538512df93648b4bb && git checkout -b ollama_vulkan_stable && \ git config user.name "Builder" && git config user.email "builder@local" && \ git remote add ollama_vanilla https://github.com/ollama/ollama.git && \ - git fetch ollama_vanilla --tags && git checkout v0.5.13 && git checkout -b ollama_vanilla_stable && \ + git fetch ollama_vanilla --tags && git checkout v0.5.14-rc0 && git checkout -b ollama_vanilla_stable && \ git checkout ollama_vulkan_stable && git merge ollama_vanilla_stable --allow-unrelated-histories --no-edit && \ - for p in patches/*.patch; do patch -p1 < $p; done + for p in /tmp/patches/00-fix-vulkan-building.patch; do patch -p1 < $p; done -# Build Shared Libraries (CPU & Vulkan) -WORKDIR /opt/ollama-vulkan -RUN cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -RUN cmake --build build --parallel -RUN cmake --install build --component CPU --strip -RUN cmake --install build --component Vulkan --strip +RUN \ + cd "/tmp/ollama-vulkan-git" && \ + make -f Makefile.sync clean sync -# Install rocm -RUN apt update -RUN apt install -y wget "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)" -RUN apt install -y python3-setuptools python3-wheel -RUN wget https://repo.radeon.com/amdgpu-install/6.3.3/ubuntu/noble/amdgpu-install_6.3.60303-1_all.deb -O /tmp/amdgpu-install_6.3.60303-1_all.deb -RUN apt install -y /tmp/amdgpu-install_6.3.60303-1_all.deb -RUN apt update && apt install -y rocm -# Build Final Binary -RUN cd /opt/ollama-vulkan && \ +FROM builder AS cpu-build +RUN \ + cd "/tmp/ollama-vulkan-git" && \ + cmake --preset CPU && cmake --build --parallel --preset CPU && \ + cmake --install build --component CPU --strip + +FROM builder AS vulkan-build +RUN \ + cd "/tmp/ollama-vulkan-git" && \ + cmake --preset Vulkan && \ + cmake --build --parallel --preset Vulkan && \ + cmake --install build --component Vulkan --strip + +FROM builder AS binary-build +RUN \ + cd "/tmp/ollama-vulkan-git" && \ . scripts/env.sh && \ mkdir -p dist/bin && \ go build -trimpath -buildmode=pie -o dist/bin/ollama . -# Final Image + FROM --platform=linux/amd64 library/ubuntu:noble -RUN apt-get update && apt-get install -y ca-certificates libcap2 libvulkan1 && \ +RUN \ + apt-get update && apt -y dist-upgrade && \ + apt-get install -y ca-certificates libcap2 libvulkan1 && \ apt-get clean && rm -rf /var/lib/apt/lists/* -# Copy Built Components -COPY --from=builder /opt/ollama-vulkan/dist/bin/ollama /bin/ollama +# Install ROCm +RUN \ + apt update && \ + apt install -y wget python3-setuptools python3-wheel && \ + wget https://repo.radeon.com/amdgpu-install/6.3.3/ubuntu/noble/amdgpu-install_6.3.60303-1_all.deb -O /tmp/amdgpu-install_6.3.60303-1_all.deb && \ + apt install -y /tmp/amdgpu-install_6.3.60303-1_all.deb && \ + apt update && apt install -y rocm && \ + apt-get clean && rm -rf /var/lib/apt/lists/* + + +COPY --from=cpu-build /tmp/ollama-vulkan-git/dist/lib/ollama/ /lib/ollama/ +COPY --from=vulkan-build /tmp/ollama-vulkan-git/dist/lib/ollama/vulkan/ /lib/ollama/vulkan/ +COPY --from=binary-build /tmp/ollama-vulkan-git/dist/bin/ /bin/ + +RUN find /lib/ollama && find /bin/ollama -# Expose Ollama Server Port EXPOSE 11434 ENV OLLAMA_HOST 0.0.0.0 -# Run Ollama Server ENTRYPOINT ["/bin/ollama"] CMD ["serve"] diff --git a/patches/00-fix-vulkan-building.patch b/patches/00-fix-vulkan-building.patch new file mode 100644 index 000000000..52e498ee2 --- /dev/null +++ b/patches/00-fix-vulkan-building.patch @@ -0,0 +1,15297 @@ +From 7c5f98c4cbfaf472a0d05baa3cc61afdcaeee7de Mon Sep 17 00:00:00 2001 +From: dream +Date: Thu, 13 Feb 2025 18:58:59 +0800 +Subject: [PATCH 2/2] fix: fix vulkan building + +1. Add preset for vulkan. +2. Add backend ggml-vulkan. +3. Add some log info. +--- + CMakePresets.json | 13 +- + discover/gpu.go | 7 +- + .../ggml/ggml/src/ggml-vulkan/CMakeLists.txt | 92 + + .../ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp | 8745 +++++++++++++++++ + .../ggml-vulkan/vulkan-shaders/CMakeLists.txt | 9 + + .../src/ggml-vulkan/vulkan-shaders/acc.comp | 29 + + .../src/ggml-vulkan/vulkan-shaders/add.comp | 29 + + .../ggml-vulkan/vulkan-shaders/argsort.comp | 69 + + .../src/ggml-vulkan/vulkan-shaders/clamp.comp | 17 + + .../ggml-vulkan/vulkan-shaders/concat.comp | 41 + + .../vulkan-shaders/contig_copy.comp | 42 + + .../src/ggml-vulkan/vulkan-shaders/copy.comp | 20 + + .../src/ggml-vulkan/vulkan-shaders/cos.comp | 17 + + .../vulkan-shaders/dequant_f32.comp | 20 + + .../vulkan-shaders/dequant_funcs.comp | 118 + + .../vulkan-shaders/dequant_funcs_cm2.comp | 325 + + .../vulkan-shaders/dequant_head.comp | 13 + + .../vulkan-shaders/dequant_iq4_nl.comp | 32 + + .../vulkan-shaders/dequant_q2_k.comp | 34 + + .../vulkan-shaders/dequant_q3_k.comp | 42 + + .../vulkan-shaders/dequant_q4_0.comp | 30 + + .../vulkan-shaders/dequant_q4_1.comp | 32 + + .../vulkan-shaders/dequant_q4_k.comp | 68 + + .../vulkan-shaders/dequant_q5_0.comp | 34 + + .../vulkan-shaders/dequant_q5_1.comp | 35 + + .../vulkan-shaders/dequant_q5_k.comp | 70 + + .../vulkan-shaders/dequant_q6_k.comp | 33 + + .../vulkan-shaders/dequant_q8_0.comp | 31 + + .../vulkan-shaders/diag_mask_inf.comp | 34 + + .../src/ggml-vulkan/vulkan-shaders/div.comp | 27 + + .../vulkan-shaders/flash_attn_cm2.comp | 289 + + .../src/ggml-vulkan/vulkan-shaders/gelu.comp | 25 + + .../vulkan-shaders/gelu_quick.comp | 23 + + .../vulkan-shaders/generic_binary_head.comp | 64 + + .../vulkan-shaders/generic_head.comp | 9 + + .../vulkan-shaders/generic_unary_head.comp | 56 + + .../ggml-vulkan/vulkan-shaders/get_rows.comp | 28 + + .../vulkan-shaders/get_rows_quant.comp | 39 + + .../vulkan-shaders/group_norm.comp | 66 + + .../ggml-vulkan/vulkan-shaders/im2col.comp | 87 + + .../vulkan-shaders/leaky_relu.comp | 22 + + .../src/ggml-vulkan/vulkan-shaders/mul.comp | 27 + + .../mul_mat_split_k_reduce.comp | 48 + + .../vulkan-shaders/mul_mat_vec.comp | 152 + + .../vulkan-shaders/mul_mat_vec_base.comp | 118 + + .../vulkan-shaders/mul_mat_vec_nc.comp | 71 + + .../vulkan-shaders/mul_mat_vec_p021.comp | 73 + + .../vulkan-shaders/mul_mat_vec_q2_k.comp | 115 + + .../vulkan-shaders/mul_mat_vec_q3_k.comp | 103 + + .../vulkan-shaders/mul_mat_vec_q4_k.comp | 133 + + .../vulkan-shaders/mul_mat_vec_q5_k.comp | 162 + + .../vulkan-shaders/mul_mat_vec_q6_k.comp | 112 + + .../ggml-vulkan/vulkan-shaders/mul_mm.comp | 631 ++ + .../vulkan-shaders/mul_mm_cm2.comp | 328 + + .../src/ggml-vulkan/vulkan-shaders/norm.comp | 44 + + .../src/ggml-vulkan/vulkan-shaders/pad.comp | 28 + + .../ggml-vulkan/vulkan-shaders/pool2d.comp | 74 + + .../src/ggml-vulkan/vulkan-shaders/relu.comp | 21 + + .../ggml-vulkan/vulkan-shaders/repeat.comp | 26 + + .../ggml-vulkan/vulkan-shaders/rms_norm.comp | 42 + + .../ggml-vulkan/vulkan-shaders/rope_head.comp | 49 + + .../ggml-vulkan/vulkan-shaders/rope_neox.comp | 37 + + .../ggml-vulkan/vulkan-shaders/rope_norm.comp | 37 + + .../src/ggml-vulkan/vulkan-shaders/scale.comp | 24 + + .../src/ggml-vulkan/vulkan-shaders/silu.comp | 22 + + .../src/ggml-vulkan/vulkan-shaders/sin.comp | 17 + + .../ggml-vulkan/vulkan-shaders/soft_max.comp | 174 + + .../ggml-vulkan/vulkan-shaders/square.comp | 17 + + .../ggml-vulkan/vulkan-shaders/sum_rows.comp | 37 + + .../src/ggml-vulkan/vulkan-shaders/tanh.comp | 20 + + .../vulkan-shaders/test_coopmat2_support.comp | 7 + + .../vulkan-shaders/timestep_embedding.comp | 41 + + .../src/ggml-vulkan/vulkan-shaders/types.comp | 323 + + .../ggml-vulkan/vulkan-shaders/upscale.comp | 36 + + .../vulkan-shaders/vulkan-shaders-gen.cpp | 594 ++ + .../src/ggml-vulkan/vulkan-shaders/wkv6.comp | 87 + + 76 files changed, 14642 insertions(+), 4 deletions(-) + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp + create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp + +diff --git a/CMakePresets.json b/CMakePresets.json +index 3ecb0a8f..a77f15ba 100644 +--- a/CMakePresets.json ++++ b/CMakePresets.json +@@ -58,7 +58,11 @@ + "cacheVariables": { + "AMDGPU_TARGETS": "gfx803;gfx900;gfx940;gfx941;gfx942;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102;gfx906:xnack-;gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-" + } +- } ++ }, ++ { ++ "name": "Vulkan", ++ "inherits": [ "Default" ] ++ } + ], + "buildPresets": [ + { +@@ -105,6 +109,11 @@ + "name": "ROCm 6", + "inherits": [ "ROCm" ], + "configurePreset": "ROCm 6" +- } ++ }, ++ { ++ "name": "Vulkan", ++ "targets": [ "ggml-vulkan" ], ++ "configurePreset": "Vulkan" ++ } + ] + } +diff --git a/discover/gpu.go b/discover/gpu.go +index ec96f5d4..8079be99 100644 +--- a/discover/gpu.go ++++ b/discover/gpu.go +@@ -197,7 +197,10 @@ func initVulkanHandles() *vulkanHandles { + libcapPaths := FindLibCapLibs() + + if len(vulkanPaths) > 0 && len(libcapPaths) > 0 { ++ slog.Info("vulkan: load libvulkan and libcap ok") + vHandles.deviceCount, vHandles.vulkan, vulkanLibPath, libcapLibPath = LoadVulkanMgmt(vulkanPaths, libcapPaths) ++ } else { ++ slog.Info("vulkan: failed to load libvulkan or libcap") + } + + return vHandles +@@ -426,7 +429,7 @@ func GetGPUInfo() GpuInfoList { + gpuInfo.ID = C.GoString(&memInfo.gpu_id[0]) + gpuInfo.Compute = fmt.Sprintf("%d.%d", memInfo.major, memInfo.minor) + gpuInfo.MinimumMemory = 0 +- gpuInfo.DependencyPath = depPaths ++ gpuInfo.DependencyPath = []string{LibOllamaPath} + gpuInfo.Name = C.GoString(&memInfo.gpu_name[0]) + gpuInfo.DriverMajor = int(memInfo.major) + gpuInfo.DriverMinor = int(memInfo.minor) +@@ -768,7 +771,7 @@ func LoadVulkanMgmt(vulkanLibPaths []string, capLibPaths []string) (int, *C.vk_h + + C.vk_init(vkLib, capLib, &resp) + if resp.err != nil { +- slog.Debug("Unable to load vulkan", "library", vkLibPath, capLibPath, "error", C.GoString(resp.err)) ++ slog.Error("Unable to load vulkan", "library", vkLibPath, capLibPath, "error", C.GoString(resp.err)) + C.free(unsafe.Pointer(resp.err)) + } else { + return int(resp.num_devices), &resp.ch, vkLibPath, capLibPath +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt b/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt +new file mode 100644 +index 00000000..9501de73 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt +@@ -0,0 +1,92 @@ ++find_package(Vulkan COMPONENTS glslc REQUIRED) ++ ++if (Vulkan_FOUND) ++ message(STATUS "Vulkan found") ++ ++ ggml_add_backend_library(ggml-vulkan ++ ggml-vulkan.cpp ++ ../../include/ggml-vulkan.h ++ ) ++ ++ # Compile a test shader to determine whether GL_NV_cooperative_matrix2 is supported. ++ # If it's not, there will be an error to stderr. ++ # If it's supported, set a define to indicate that we should compile those shaders ++ execute_process(COMMAND ${Vulkan_GLSLC_EXECUTABLE} -o - -fshader-stage=compute --target-env=vulkan1.3 "${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders/test_coopmat2_support.comp" ++ OUTPUT_VARIABLE glslc_output ++ ERROR_VARIABLE glslc_error) ++ ++ if (${glslc_error} MATCHES ".*extension not supported: GL_NV_cooperative_matrix2.*") ++ message(STATUS "GL_NV_cooperative_matrix2 not supported by glslc") ++ else() ++ message(STATUS "GL_NV_cooperative_matrix2 supported by glslc") ++ add_compile_definitions(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) ++ endif() ++ ++ target_link_libraries(ggml-vulkan PRIVATE Vulkan::Vulkan) ++ target_include_directories(ggml-vulkan PRIVATE ${CMAKE_CURRENT_BINARY_DIR}) ++ ++ # Workaround to the "can't dereference invalidated vector iterator" bug in clang-cl debug build ++ # Posssibly relevant: https://stackoverflow.com/questions/74748276/visual-studio-no-displays-the-correct-length-of-stdvector ++ if (MSVC AND CMAKE_CXX_COMPILER_ID STREQUAL "Clang") ++ add_compile_definitions(_ITERATOR_DEBUG_LEVEL=0) ++ endif() ++ ++ if (GGML_VULKAN_CHECK_RESULTS) ++ add_compile_definitions(GGML_VULKAN_CHECK_RESULTS) ++ endif() ++ ++ if (GGML_VULKAN_DEBUG) ++ add_compile_definitions(GGML_VULKAN_DEBUG) ++ endif() ++ ++ if (GGML_VULKAN_MEMORY_DEBUG) ++ add_compile_definitions(GGML_VULKAN_MEMORY_DEBUG) ++ endif() ++ ++ if (GGML_VULKAN_SHADER_DEBUG_INFO) ++ add_compile_definitions(GGML_VULKAN_SHADER_DEBUG_INFO) ++ endif() ++ ++ if (GGML_VULKAN_PERF) ++ add_compile_definitions(GGML_VULKAN_PERF) ++ endif() ++ ++ if (GGML_VULKAN_VALIDATE) ++ add_compile_definitions(GGML_VULKAN_VALIDATE) ++ endif() ++ ++ if (GGML_VULKAN_RUN_TESTS) ++ add_compile_definitions(GGML_VULKAN_RUN_TESTS) ++ endif() ++ ++ add_subdirectory(vulkan-shaders) ++ ++ set (_ggml_vk_genshaders_cmd vulkan-shaders-gen) ++ set (_ggml_vk_header ${CMAKE_CURRENT_BINARY_DIR}/ggml-vulkan-shaders.hpp) ++ set (_ggml_vk_source ${CMAKE_CURRENT_BINARY_DIR}/ggml-vulkan-shaders.cpp) ++ set (_ggml_vk_input_dir ${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders) ++ set (_ggml_vk_output_dir ${CMAKE_CURRENT_BINARY_DIR}/vulkan-shaders.spv) ++ ++ file(GLOB _ggml_vk_shader_deps "${_ggml_vk_input_dir}/*.comp") ++ ++ add_custom_command( ++ OUTPUT ${_ggml_vk_header} ++ ${_ggml_vk_source} ++ ++ COMMAND "$/${_ggml_vk_genshaders_cmd}" ++ --glslc ${Vulkan_GLSLC_EXECUTABLE} ++ --input-dir ${_ggml_vk_input_dir} ++ --output-dir ${_ggml_vk_output_dir} ++ --target-hpp ${_ggml_vk_header} ++ --target-cpp ${_ggml_vk_source} ++ --no-clean ++ ++ DEPENDS ${_ggml_vk_shader_deps} ${_ggml_vk_genshaders_cmd} ++ COMMENT "Generate vulkan shaders" ++ ) ++ ++ target_sources(ggml-vulkan PRIVATE ${_ggml_vk_source} ${_ggml_vk_header}) ++ ++else() ++ message(WARNING "Vulkan not found") ++endif() +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp +new file mode 100644 +index 00000000..d75cd6d6 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp +@@ -0,0 +1,8745 @@ ++#include "ggml-vulkan.h" ++#include ++#if defined(GGML_VULKAN_RUN_TESTS) || defined(GGML_VULKAN_PERF) || defined(GGML_VULKAN_CHECK_RESULTS) ++#include ++#include "ggml-cpu.h" ++#endif ++ ++#include ++ ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++ ++#include "ggml-impl.h" ++#include "ggml-backend-impl.h" ++ ++#include "ggml-vulkan-shaders.hpp" ++ ++#define VK_API_VERSION VK_API_VERSION_1_2 ++ ++#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) ++ ++#define VK_VENDOR_ID_AMD 0x1002 ++#define VK_VENDOR_ID_APPLE 0x106b ++#define VK_VENDOR_ID_INTEL 0x8086 ++#define VK_VENDOR_ID_NVIDIA 0x10de ++ ++#define VK_DEVICE_DESCRIPTOR_POOL_SIZE 32 ++ ++#define GGML_VK_MAX_NODES 8192 ++ ++#define MAX_VK_BUFFERS 256 ++ ++#define VK_CHECK(err, msg) \ ++ do { \ ++ vk::Result err_ = (err); \ ++ if (err_ != vk::Result::eSuccess) { \ ++ fprintf(stderr, "ggml_vulkan: %s error %s at %s:%d\n", \ ++ #err, to_string(err_).c_str(), __FILE__, __LINE__); \ ++ exit(1); \ ++ } \ ++ } while (0) ++ ++#ifdef GGML_VULKAN_DEBUG ++#define VK_LOG_DEBUG(msg) std::cerr << msg << std::endl ++#else ++#define VK_LOG_DEBUG(msg) ((void) 0) ++#endif // GGML_VULKAN_DEBUG ++ ++struct ggml_backend_vk_context; ++ ++struct vk_queue { ++ uint32_t queue_family_index; ++ vk::Queue queue; ++ vk::CommandPool pool; ++ uint32_t cmd_buffer_idx; ++ std::vector cmd_buffers; ++ ++ vk::PipelineStageFlags stage_flags; ++ ++ bool transfer_only; ++}; ++ ++struct vk_pipeline_struct { ++ std::string name; ++ vk::ShaderModule shader_module; ++ vk::DescriptorSetLayout dsl; ++ std::vector descriptor_pools; ++ std::vector descriptor_sets; ++ uint32_t descriptor_set_idx; ++ vk::PipelineLayout layout; ++ vk::Pipeline pipeline; ++ uint32_t push_constant_size; ++ uint32_t parameter_count; ++ std::array wg_denoms; ++ uint32_t align; ++}; ++ ++typedef std::shared_ptr vk_pipeline; ++typedef std::weak_ptr vk_pipeline_ref; ++ ++static void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline); ++ ++struct vk_matmul_pipeline_struct { ++ vk_pipeline l, m, s; ++ vk_pipeline a_l, a_m, a_s; ++}; ++ ++typedef std::shared_ptr vk_matmul_pipeline; ++ ++struct vk_matmul_pipeline2 { ++ vk_matmul_pipeline2() { ++ f16acc = std::make_shared(); ++ f32acc = std::make_shared(); ++ } ++ vk_matmul_pipeline f32acc; ++ vk_matmul_pipeline f16acc; ++}; ++ ++struct vk_device_struct; ++typedef std::shared_ptr vk_device; ++typedef std::weak_ptr vk_device_ref; ++ ++struct vk_buffer_struct; ++typedef std::shared_ptr vk_buffer; ++typedef std::weak_ptr vk_buffer_ref; ++ ++struct ggml_backend_vk_buffer_type_context { ++ std::string name; ++ vk_device device; ++}; ++ ++static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft); ++static ggml_backend_buffer_t ggml_backend_vk_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size); ++static size_t ggml_backend_vk_buffer_type_get_alignment(ggml_backend_buffer_type_t buft); ++static size_t ggml_backend_vk_buffer_type_get_max_size(ggml_backend_buffer_type_t buft); ++static size_t ggml_backend_vk_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor); ++static ggml_backend_buffer_type_i ggml_backend_vk_buffer_type_interface = { ++ /* .get_name = */ ggml_backend_vk_buffer_type_name, ++ /* .alloc_buffer = */ ggml_backend_vk_buffer_type_alloc_buffer, ++ /* .get_alignment = */ ggml_backend_vk_buffer_type_get_alignment, ++ /* .get_max_size = */ ggml_backend_vk_buffer_type_get_max_size, ++ /* .get_alloc_size = */ ggml_backend_vk_buffer_type_get_alloc_size, ++ /* .is_host = */ NULL, ++}; ++ ++#ifdef GGML_VULKAN_MEMORY_DEBUG ++class vk_memory_logger; ++#endif ++#ifdef GGML_VULKAN_PERF ++class vk_perf_logger; ++#endif ++static void ggml_vk_destroy_buffer(vk_buffer& buf); ++ ++static constexpr uint32_t mul_mat_vec_max_cols = 8; ++ ++struct vk_device_struct { ++ std::mutex mutex; ++ ++ vk::PhysicalDevice physical_device; ++ vk::PhysicalDeviceProperties properties; ++ std::string name; ++ uint64_t max_memory_allocation_size; ++ bool fp16; ++ bool pipeline_robustness; ++ vk::Device device; ++ uint32_t vendor_id; ++ vk_queue compute_queue; ++ vk_queue transfer_queue; ++ bool single_queue; ++ uint32_t subgroup_size; ++ uint32_t shader_core_count; ++ bool uma; ++ bool float_controls_rte_fp16; ++ ++ bool subgroup_size_control; ++ uint32_t subgroup_min_size; ++ uint32_t subgroup_max_size; ++ bool subgroup_require_full_support; ++ ++ bool coopmat_support; ++ bool coopmat_acc_f32_support; ++ bool coopmat_acc_f16_support; ++ uint32_t coopmat_m; ++ uint32_t coopmat_n; ++ uint32_t coopmat_k; ++ bool coopmat2; ++ ++ size_t idx; ++ ++ bool mul_mat_l; ++ bool mul_mat_m; ++ bool mul_mat_s; ++ bool mul_mat_id_l; ++ bool mul_mat_id_m; ++ bool mul_mat_id_s; ++ ++ vk_matmul_pipeline pipeline_matmul_f32; ++ vk_matmul_pipeline pipeline_matmul_f32_f16; ++ vk_matmul_pipeline2 pipeline_matmul_f16; ++ vk_matmul_pipeline2 pipeline_matmul_f16_f32; ++ vk_pipeline pipeline_matmul_split_k_reduce; ++ ++ vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_COUNT]; ++ vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat[GGML_TYPE_COUNT]; ++ ++ vk_matmul_pipeline pipeline_matmul_id_f32; ++ vk_matmul_pipeline2 pipeline_matmul_id_f16; ++ vk_matmul_pipeline2 pipeline_matmul_id_f16_f32; ++ ++ vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id[GGML_TYPE_COUNT]; ++ ++ vk_pipeline pipeline_dequant[GGML_TYPE_COUNT]; ++ vk_pipeline pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_COUNT][mul_mat_vec_max_cols]; ++ vk_pipeline pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_COUNT][mul_mat_vec_max_cols]; ++ vk_pipeline pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_COUNT]; ++ ++ vk_pipeline pipeline_mul_mat_vec_p021_f16_f32; ++ vk_pipeline pipeline_mul_mat_vec_nc_f16_f32; ++ vk_pipeline pipeline_get_rows[GGML_TYPE_COUNT]; ++ vk_pipeline pipeline_get_rows_f32[GGML_TYPE_COUNT]; ++ vk_pipeline pipeline_acc_f32; ++ vk_pipeline pipeline_add_f32, pipeline_add_f32_norepeat; ++ vk_pipeline pipeline_add_f16_f32_f16, pipeline_add_f16_f32_f16_norepeat; ++ vk_pipeline pipeline_mul_f32, pipeline_mul_f32_norepeat; ++ vk_pipeline pipeline_div_f32, pipeline_div_f32_norepeat; ++ vk_pipeline pipeline_concat_f32, pipeline_concat_f16, pipeline_concat_i32; ++ vk_pipeline pipeline_upscale_f32; ++ vk_pipeline pipeline_scale_f32; ++ vk_pipeline pipeline_sqr_f32; ++ vk_pipeline pipeline_sin_f32; ++ vk_pipeline pipeline_cos_f32; ++ vk_pipeline pipeline_clamp_f32; ++ vk_pipeline pipeline_pad_f32; ++ vk_pipeline pipeline_repeat_f32; ++ vk_pipeline pipeline_cpy_f32_f32, pipeline_cpy_f32_f16, pipeline_cpy_f16_f16; ++ vk_pipeline pipeline_contig_cpy_f32_f32, pipeline_contig_cpy_f32_f16, pipeline_contig_cpy_f16_f16; ++ vk_pipeline pipeline_norm_f32; ++ vk_pipeline pipeline_group_norm_f32; ++ vk_pipeline pipeline_rms_norm_f32; ++ vk_pipeline pipeline_gelu_f32; ++ vk_pipeline pipeline_gelu_quick_f32; ++ vk_pipeline pipeline_silu_f32; ++ vk_pipeline pipeline_relu_f32; ++ vk_pipeline pipeline_leaky_relu_f32; ++ vk_pipeline pipeline_tanh_f32; ++ vk_pipeline pipeline_diag_mask_inf_f32; ++ vk_pipeline pipeline_soft_max_f32, pipeline_soft_max_f32_f16; ++ vk_pipeline pipeline_soft_max_f32_wg512, pipeline_soft_max_f32_f16_wg512; ++ vk_pipeline pipeline_rope_norm_f32, pipeline_rope_norm_f16; ++ vk_pipeline pipeline_rope_neox_f32, pipeline_rope_neox_f16; ++ vk_pipeline pipeline_argsort_f32; ++ vk_pipeline pipeline_sum_rows_f32; ++ vk_pipeline pipeline_im2col_f32, pipeline_im2col_f32_f16; ++ vk_pipeline pipeline_timestep_embedding_f32; ++ vk_pipeline pipeline_pool2d_f32; ++ vk_pipeline pipeline_rwkv_wkv6_f32; ++ ++ // [2][2][2] is for {f16acc,f32acc}x{large,small_rows}x{unaligned, aligned} ++ vk_pipeline pipeline_flash_attn_f32_f16_D64[GGML_TYPE_COUNT][2][2][2]; ++ vk_pipeline pipeline_flash_attn_f32_f16_D80[GGML_TYPE_COUNT][2][2][2]; ++ vk_pipeline pipeline_flash_attn_f32_f16_D96[GGML_TYPE_COUNT][2][2][2]; ++ vk_pipeline pipeline_flash_attn_f32_f16_D112[GGML_TYPE_COUNT][2][2][2]; ++ vk_pipeline pipeline_flash_attn_f32_f16_D128[GGML_TYPE_COUNT][2][2][2]; ++ vk_pipeline pipeline_flash_attn_f32_f16_D256[GGML_TYPE_COUNT][2][2][2]; ++ ++ std::unordered_map pipelines; ++ std::unordered_map pipeline_descriptor_set_requirements; ++ ++ std::vector> pinned_memory; ++ ++ vk::Fence fence; ++ vk_buffer sync_staging; ++ ++ ggml_backend_buffer_type buffer_type; ++ ++#ifdef GGML_VULKAN_MEMORY_DEBUG ++ std::unique_ptr memory_logger; ++#endif ++#ifdef GGML_VULKAN_PERF ++ std::unique_ptr perf_logger; ++#endif ++ ++ ~vk_device_struct() { ++ VK_LOG_DEBUG("destroy device " << name); ++ ++ device.destroyFence(fence); ++ ++ ggml_vk_destroy_buffer(sync_staging); ++ ++ device.destroyCommandPool(compute_queue.pool); ++ if (!single_queue) { ++ device.destroyCommandPool(transfer_queue.pool); ++ } ++ ++ for (auto& pipeline : pipelines) { ++ if (pipeline.second.expired()) { ++ continue; ++ } ++ ++ vk_pipeline pl = pipeline.second.lock(); ++ ggml_vk_destroy_pipeline(device, pl); ++ } ++ pipelines.clear(); ++ ++ device.destroy(); ++ } ++}; ++ ++struct vk_buffer_struct { ++ vk::Buffer buffer = VK_NULL_HANDLE; ++ vk::DeviceMemory device_memory = VK_NULL_HANDLE; ++ vk::MemoryPropertyFlags memory_property_flags; ++ void * ptr; ++ size_t size = 0; ++ ++ vk_device device; ++ ++ ~vk_buffer_struct() { ++ if (size == 0) { ++ return; ++ } ++ VK_LOG_DEBUG("~vk_buffer_struct(" << buffer << ", " << size << ")"); ++ ++ device->device.freeMemory(device_memory); ++ device->device.destroyBuffer(buffer); ++ } ++}; ++ ++struct vk_subbuffer { ++ vk_buffer buffer; ++ uint64_t offset; ++ uint64_t size; ++ ++ operator vk::DescriptorBufferInfo() const { ++ return { buffer->buffer, offset, size }; ++ } ++}; ++ ++struct vk_semaphore { ++ vk::Semaphore s; ++ uint64_t value; ++}; ++ ++struct vk_submission { ++ vk::CommandBuffer buffer; ++ std::vector wait_semaphores; ++ std::vector signal_semaphores; ++}; ++ ++typedef std::vector vk_sequence; ++ ++struct vk_mat_mat_push_constants { ++ uint32_t M; uint32_t N; uint32_t K; ++ uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; ++ uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; ++ uint32_t k_split; ++ uint32_t ne02; uint32_t ne12; uint32_t broadcast2; uint32_t broadcast3; ++}; ++struct vk_mat_vec_push_constants { ++ uint32_t ncols; uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; ++ uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; ++ uint32_t ne02; uint32_t ne12; uint32_t broadcast2; uint32_t broadcast3; ++}; ++ ++struct vk_mat_mat_id_push_constants { ++ uint32_t M; uint32_t N; uint32_t K; ++ uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; ++ uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; ++ uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11; ++}; ++struct vk_mat_vec_id_push_constants { ++ uint32_t ncols; uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; ++ uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; ++ uint32_t nei0; uint32_t ne11; ++}; ++ ++struct vk_flash_attn_push_constants { ++ uint32_t N; ++ uint32_t KV; ++ ++ uint32_t ne1; ++ uint32_t ne2; ++ uint32_t ne3; ++ ++ uint32_t neq2; ++ uint32_t neq3; ++ uint32_t nek2; ++ uint32_t nek3; ++ uint32_t nev2; ++ uint32_t nev3; ++ uint32_t nem1; ++ ++ uint32_t nb02; ++ uint32_t nb03; ++ uint32_t nb12; ++ uint32_t nb13; ++ uint32_t nb22; ++ uint32_t nb23; ++ uint32_t nb31; ++ ++ float scale; ++ float max_bias; ++ float logit_softcap; ++ ++ uint32_t mask; ++ uint32_t n_head_log2; ++ float m0; ++ float m1; ++}; ++ ++struct vk_op_push_constants { ++ uint32_t KX; ++ uint32_t KY; ++ float param1; ++ float param2; ++}; ++ ++struct vk_op_unary_push_constants { ++ uint32_t ne; ++ uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; ++ uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; uint32_t nb10; uint32_t nb11; uint32_t nb12; uint32_t nb13; ++ uint32_t misalign_offsets; ++ float param1; float param2; ++ uint32_t ne0_012mp; uint32_t ne0_012L; ++ uint32_t ne0_01mp; uint32_t ne0_01L; ++ uint32_t ne0_0mp; uint32_t ne0_0L; ++ uint32_t ne1_012mp; uint32_t ne1_012L; ++ uint32_t ne1_01mp; uint32_t ne1_01L; ++ uint32_t ne1_0mp; uint32_t ne1_0L; ++}; ++static_assert(sizeof(vk_op_unary_push_constants) <= 128, "sizeof(vk_op_unary_push_constants) must be <= 128"); ++ ++// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1. ++// Precompute mp (m' in the paper) and L such that division ++// can be computed using a multiply (high 32b of 64b result) ++// and a shift: ++// ++// n/d = (mulhi(n, mp) + n) >> L; ++static void init_fastdiv_values(uint32_t d, uint32_t &mp, uint32_t &L) ++{ ++ // compute L = ceil(log2(d)); ++ L = 0; ++ while (L < 32 && (uint32_t{1} << L) < d) { ++ L++; ++ } ++ ++ mp = (uint32_t)((uint64_t{1} << 32) * ((uint64_t{1} << L) - d) / d + 1); ++} ++ ++template void init_pushconst_fastdiv(T &p) { ++ GGML_UNUSED(p); ++ static_assert(!std::is_const::value, "unexpected type"); ++} ++ ++template <> void init_pushconst_fastdiv(vk_op_unary_push_constants &p) { ++ // Compute magic values to divide by these six numbers. ++ init_fastdiv_values(p.ne02*p.ne01*p.ne00, p.ne0_012mp, p.ne0_012L); ++ init_fastdiv_values(p.ne01*p.ne00, p.ne0_01mp, p.ne0_01L); ++ init_fastdiv_values(p.ne00, p.ne0_0mp, p.ne0_0L); ++ init_fastdiv_values(p.ne12*p.ne11*p.ne10, p.ne1_012mp, p.ne1_012L); ++ init_fastdiv_values(p.ne11*p.ne10, p.ne1_01mp, p.ne1_01L); ++ init_fastdiv_values(p.ne10, p.ne1_0mp, p.ne1_0L); ++} ++ ++struct vk_op_binary_push_constants { ++ uint32_t ne; ++ uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; ++ uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; uint32_t nb10; uint32_t nb11; uint32_t nb12; uint32_t nb13; ++ uint32_t ne20; uint32_t ne21; uint32_t ne22; uint32_t ne23; uint32_t nb20; uint32_t nb21; uint32_t nb22; uint32_t nb23; ++ uint32_t misalign_offsets; ++ float param1; float param2; int32_t param3; ++}; ++ ++struct vk_op_diag_mask_push_constants { ++ uint32_t ncols; ++ uint32_t rows_per_channel; ++ int32_t n_past; ++}; ++ ++struct vk_op_rope_push_constants { ++ uint32_t ncols; ++ uint32_t n_dims; ++ float freq_scale; ++ uint32_t p_delta_rows; ++ float freq_base; ++ float ext_factor; ++ float attn_factor; ++ float corr_dims[2]; ++ float theta_scale; ++ uint32_t has_ff; ++}; ++ ++struct vk_op_soft_max_push_constants { ++ uint32_t KX; ++ uint32_t KY; ++ float scale; ++ float max_bias; ++ float m0; ++ float m1; ++ uint32_t n_head_log2; ++ uint32_t nrows_x; ++}; ++ ++struct vk_op_argsort_push_constants { ++ uint32_t ncols; ++ uint32_t ncols_pad; ++ int32_t order; ++}; ++ ++struct vk_op_im2col_push_constants { ++ uint32_t batch_offset; uint32_t offset_delta; ++ uint32_t IC; ++ uint32_t IW; uint32_t IH; ++ uint32_t OW; uint32_t OH; ++ uint32_t KW; uint32_t KH; ++ uint32_t pelements; ++ uint32_t CHW; ++ int32_t s0; int32_t s1; ++ int32_t p0; int32_t p1; ++ int32_t d0; int32_t d1; ++}; ++ ++struct vk_op_timestep_embedding_push_constants { ++ uint32_t nb1; ++ uint32_t dim; ++ uint32_t max_period; ++}; ++ ++struct vk_op_pool2d_push_constants { ++ uint32_t IW; uint32_t IH; ++ uint32_t OW; uint32_t OH; ++ uint32_t OC; ++ uint32_t pelements; ++ uint32_t op; ++ int32_t k0; int32_t k1; ++ int32_t s0; int32_t s1; ++ int32_t p0; int32_t p1; ++}; ++ ++struct vk_op_rwkv_wkv6_push_constants { ++ uint32_t B; ++ uint32_t T; ++ uint32_t C; ++ uint32_t H; ++}; ++ ++// Allow pre-recording command buffers ++struct vk_staging_memcpy { ++ vk_staging_memcpy(void * _dst, const void * _src, size_t _n) : dst(_dst), src(_src), n(_n) {} ++ ++ void * dst; ++ const void * src; ++ size_t n; ++}; ++ ++struct vk_op_upscale_push_constants { ++ uint32_t ne; uint32_t a_offset; uint32_t d_offset; ++ uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; ++ uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; ++ float sf0; float sf1; float sf2; float sf3; ++}; ++ ++struct vk_context_struct { ++ vk_submission * s; ++ std::vector seqs; ++ ++ int exit_tensor_idx; ++ ++ std::vector in_memcpys; ++ std::vector out_memcpys; ++ ++ vk_queue * q; ++}; ++typedef std::shared_ptr vk_context; ++typedef std::weak_ptr vk_context_ref; ++ ++struct ggml_vk_garbage_collector { ++ std::vector tl_semaphores; ++ std::vector semaphores; ++ std::vector events; ++ std::vector temp_buffers; ++ std::vector contexts; ++}; ++ ++#if defined(GGML_VULKAN_MEMORY_DEBUG) || defined(GGML_VULKAN_DEBUG) ++#define VK_LOG_MEMORY(msg) std::cerr << "ggml_vulkan memory: " << msg << std::endl ++ ++static std::string format_size(size_t size) { ++ const size_t kib = 1024; ++ const size_t mib = kib * 1024; ++ const size_t gib = mib * 1024; ++ ++ std::ostringstream oss; ++ oss << std::fixed << std::setprecision(2); ++ ++ if (size >= gib) { ++ oss << static_cast(size) / gib << " GiB"; ++ } else if (size >= mib) { ++ oss << static_cast(size) / mib << " MiB"; ++ } else if (size >= kib) { ++ oss << static_cast(size) / kib << " KiB"; ++ } else { ++ oss << size << " B"; ++ } ++ ++ return oss.str(); ++} ++ ++static std::mutex log_mutex; ++ ++class vk_memory_logger { ++public: ++ vk_memory_logger(): total_device(0), total_host(0) {} ++ void log_allocation(vk_buffer_ref buf_ref, size_t size); ++ void log_deallocation(vk_buffer_ref buf_ref); ++ ++private: ++ std::map allocations; // Track allocations ++ size_t total_device; ++ size_t total_host; ++}; ++#else ++#define VK_LOG_MEMORY(msg) ((void) 0) ++#endif // GGML_VULKAN_MEMORY_DEBUG ++ ++#if defined(GGML_VULKAN_PERF) ++ ++class vk_perf_logger { ++public: ++ void print_timings() { ++ std::cerr << "----------------\nVulkan Timings:" << std::endl; ++ for (const auto& t : timings) { ++ uint64_t total = 0; ++ for (const auto& time : t.second) { ++ total += time; ++ } ++ std::cerr << t.first << ": " << t.second.size() << " x " << (total / t.second.size() / 1000.0) << " ms" << std::endl; ++ } ++ ++ timings.clear(); ++ } ++ ++ void log_timing(const ggml_tensor * node, uint64_t time) { ++ if (node->op == GGML_OP_UNARY) { ++ timings[ggml_unary_op_name(ggml_get_unary_op(node))].push_back(time); ++ return; ++ } ++ if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { ++ const uint64_t m = node->src[0]->ne[1]; ++ const uint64_t n = node->src[1]->ne[1]; ++ const uint64_t k = node->src[1]->ne[0]; ++ std::string name = ggml_op_name(node->op); ++ if (n == 1) { ++ name += "_VEC m=" + std::to_string(m) + " k=" + std::to_string(k); ++ } else { ++ name += " m=" + std::to_string(m) + " n=" + std::to_string(n) + " k=" + std::to_string(k); ++ } ++ timings[name].push_back(time); ++ return; ++ } ++ timings[ggml_op_name(node->op)].push_back(time); ++ } ++private: ++ std::map> timings; ++}; ++#endif // GGML_VULKAN_PERF ++ ++struct ggml_backend_vk_context { ++ std::string name; ++ ++ vk_device device; ++ ++ size_t semaphore_idx, event_idx; ++ ggml_vk_garbage_collector gc; ++ size_t prealloc_size_x, prealloc_size_y, prealloc_size_split_k; ++ vk_buffer prealloc_x, prealloc_y, prealloc_split_k; ++ vk::Fence fence; ++ ++ vk_buffer buffer_pool[MAX_VK_BUFFERS]; ++ ++ vk_context_ref compute_ctx; ++ vk_context_ref transfer_ctx; ++ ++ std::vector tensor_ctxs; ++}; ++ ++static void * const vk_ptr_base = (void *)(uintptr_t) 0x1000; // NOLINT ++ ++static uint64_t vk_tensor_offset(const ggml_tensor * tensor) { ++ if (tensor->view_src) { ++ return (uint8_t *) tensor->view_src->data - (uint8_t *) vk_ptr_base; ++ } ++ return (uint8_t *) tensor->data - (uint8_t *) vk_ptr_base; ++} ++ ++struct ggml_backend_vk_buffer_context { ++ vk_device_ref device; ++ vk_buffer dev_buffer; ++ std::string name; ++ ++ ggml_backend_vk_buffer_context(vk_device_ref device, vk_buffer&& dev_buffer, std::string& name) : ++ device(device), ++ dev_buffer(dev_buffer), ++ name(name) { ++ } ++ ++ ~ggml_backend_vk_buffer_context() { ++ ggml_vk_destroy_buffer(dev_buffer); ++ } ++}; ++ ++#ifdef GGML_VULKAN_MEMORY_DEBUG ++void vk_memory_logger::log_allocation(vk_buffer_ref buf_ref, size_t size) { ++ std::lock_guard guard(log_mutex); ++ vk_buffer buf = buf_ref.lock(); ++ const bool device = bool(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eDeviceLocal); ++ const std::string type = device ? "device" : "host"; ++ allocations[buf->buffer] = size; ++ total_device += device ? size : 0; ++ total_host += device ? 0 : size; ++ VK_LOG_MEMORY(buf->device->name << ": +" << format_size(size) << " " << type << " at " << buf->buffer << ". Total device: " << format_size(total_device) << ", total host: " << format_size(total_host)); ++} ++ ++void vk_memory_logger::log_deallocation(vk_buffer_ref buf_ref) { ++ if (buf_ref.expired() || buf_ref.lock()->size == 0) { ++ return; ++ } ++ ++ std::lock_guard guard(log_mutex); ++ vk_buffer buf = buf_ref.lock(); ++ const bool device = bool(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eDeviceLocal); ++ std::string type = device ? "device" : "host"; ++ auto it = allocations.find(buf->buffer); ++ total_device -= device ? it->second : 0; ++ total_host -= device ? 0 : it->second; ++ if (it != allocations.end()) { ++ VK_LOG_MEMORY(buf->device->name << ": -" << format_size(it->second) << " " << type << " at " << buf->buffer << ". Total device: " << format_size(total_device) << ", total host: " << format_size(total_host)); ++ allocations.erase(it); ++ } else { ++ VK_LOG_MEMORY("ERROR " << buf->device->name << ": Attempted to deallocate unknown " << type << " memory at " << buf->buffer); ++ } ++} ++#endif // GGML_VULKAN_MEMORY_DEBUG ++ ++struct vk_instance_t { ++ vk::Instance instance; ++ ++ std::vector device_indices; ++ vk_device devices[GGML_VK_MAX_DEVICES]; ++}; ++ ++static bool vk_instance_initialized = false; ++static vk_instance_t vk_instance; ++ ++#ifdef GGML_VULKAN_CHECK_RESULTS ++static size_t vk_skip_checks; ++static size_t vk_output_tensor; ++ ++static void ggml_vk_print_tensor(const ggml_tensor * tensor, const char * name); ++static void ggml_vk_check_results_0(ggml_tensor * tensor); ++static void ggml_vk_check_results_1(ggml_tensor * tensor); ++#endif ++ ++typedef void (*ggml_vk_func_t)(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); ++ ++static void ggml_backend_vk_free(ggml_backend_t backend); ++ ++// variables to track number of compiles in progress ++static uint32_t compile_count = 0; ++static std::mutex compile_count_mutex; ++static std::condition_variable compile_count_cond; ++ ++static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipeline, const std::string name, size_t spv_size, const void* spv_data, const std::string entrypoint, ++ uint32_t parameter_count, uint32_t push_constant_size, std::array wg_denoms, std::vector specialization_constants, ++ uint32_t align, bool disable_robustness, bool require_full_subgroups, uint32_t required_subgroup_size) { ++ VK_LOG_DEBUG("ggml_vk_create_pipeline(" << device->name << ", " << name << ", " << entrypoint << ", " << parameter_count << ", " << push_constant_size << ++ ", (" << wg_denoms[0] << "," << wg_denoms[1] << "," << wg_denoms[2] << "), specialization_constants, " << align << ++ ", " << disable_robustness << ", " << require_full_subgroups << ", " << required_subgroup_size << ")"); ++ GGML_ASSERT(parameter_count > 0); ++ GGML_ASSERT(wg_denoms[0] > 0 && wg_denoms[1] > 0 && wg_denoms[2] > 0); // NOLINT ++ ++ pipeline = std::make_shared(); ++ pipeline->name = name; ++ pipeline->parameter_count = parameter_count; ++ pipeline->push_constant_size = push_constant_size; ++ pipeline->wg_denoms = wg_denoms; ++ pipeline->align = align; ++ ++ vk::ShaderModuleCreateInfo shader_module_create_info({}, spv_size, reinterpret_cast(spv_data)); ++ pipeline->shader_module = device->device.createShaderModule(shader_module_create_info); ++ ++ std::vector dsl_binding; ++ std::vector dsl_binding_flags; ++ for (uint32_t i = 0; i < parameter_count; i++) { ++ dsl_binding.push_back({i, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute}); ++ dsl_binding_flags.push_back({}); ++ } ++ ++ vk::DescriptorSetLayoutBindingFlagsCreateInfo dslbfci = { dsl_binding_flags }; ++ ++ vk::PushConstantRange pcr( ++ vk::ShaderStageFlagBits::eCompute, ++ 0, ++ pipeline->push_constant_size ++ ); ++ ++ vk::DescriptorSetLayoutCreateInfo descriptor_set_layout_create_info( ++ {}, ++ dsl_binding); ++ descriptor_set_layout_create_info.setPNext(&dslbfci); ++ pipeline->dsl = device->device.createDescriptorSetLayout(descriptor_set_layout_create_info); ++ ++ vk::DescriptorPoolSize descriptor_pool_size(vk::DescriptorType::eStorageBuffer, pipeline->parameter_count * VK_DEVICE_DESCRIPTOR_POOL_SIZE); ++ vk::DescriptorPoolCreateInfo descriptor_pool_create_info({}, VK_DEVICE_DESCRIPTOR_POOL_SIZE, descriptor_pool_size); ++ pipeline->descriptor_pools.push_back(device->device.createDescriptorPool(descriptor_pool_create_info)); ++ ++ pipeline->descriptor_set_idx = 0; ++ ++ vk::PipelineLayoutCreateInfo pipeline_layout_create_info(vk::PipelineLayoutCreateFlags(), pipeline->dsl, pcr); ++ pipeline->layout = device->device.createPipelineLayout(pipeline_layout_create_info); ++ ++ std::vector specialization_entries(specialization_constants.size()); ++ ++ for (size_t i = 0; i < specialization_constants.size(); i++) { ++ specialization_entries[i].constantID = i; ++ specialization_entries[i].offset = i * sizeof(uint32_t); ++ specialization_entries[i].size = sizeof(uint32_t); ++ } ++ ++ vk::SpecializationInfo specialization_info( ++ specialization_entries.size(), ++ specialization_entries.data(), ++ specialization_constants.size() * sizeof(uint32_t), ++ specialization_constants.data() ++ ); ++ ++ vk::PipelineShaderStageCreateFlags pipeline_shader_stage_create_flags{}; ++ ++ if (device->subgroup_require_full_support && require_full_subgroups) { ++ pipeline_shader_stage_create_flags |= vk::PipelineShaderStageCreateFlagBits::eRequireFullSubgroupsEXT; ++ } ++ ++ vk::PipelineShaderStageCreateInfo pipeline_shader_create_info( ++ pipeline_shader_stage_create_flags, ++ vk::ShaderStageFlagBits::eCompute, ++ pipeline->shader_module, ++ entrypoint.c_str(), ++ &specialization_info); ++ ++ vk::PipelineShaderStageRequiredSubgroupSizeCreateInfoEXT pipeline_shader_stage_required_subgroup_size_create_info; ++ pipeline_shader_stage_required_subgroup_size_create_info.requiredSubgroupSize = required_subgroup_size; ++ if (device->subgroup_size_control && required_subgroup_size > 0) { ++ GGML_ASSERT(device->subgroup_min_size <= required_subgroup_size && required_subgroup_size <= device->subgroup_max_size); ++ pipeline_shader_create_info.setPNext(&pipeline_shader_stage_required_subgroup_size_create_info); ++ } ++ ++ vk::ComputePipelineCreateInfo compute_pipeline_create_info( ++ vk::PipelineCreateFlags{}, ++ pipeline_shader_create_info, ++ pipeline->layout); ++ ++ vk::PipelineRobustnessCreateInfoEXT rci; ++ ++ if (device->pipeline_robustness && disable_robustness) { ++ rci.storageBuffers = vk::PipelineRobustnessBufferBehaviorEXT::eDisabled; ++ rci.uniformBuffers = vk::PipelineRobustnessBufferBehaviorEXT::eDisabled; ++ compute_pipeline_create_info.setPNext(&rci); ++ } ++ ++ pipeline->pipeline = device->device.createComputePipeline(VK_NULL_HANDLE, compute_pipeline_create_info).value; ++ ++ { ++ std::lock_guard guard(device->mutex); ++ device->pipelines.insert({ pipeline->name, pipeline }); ++ } ++ ++ { ++ std::lock_guard guard(compile_count_mutex); ++ assert(compile_count > 0); ++ compile_count--; ++ ++ // "Progress bar" for shader compiles ++ static uint32_t total_compile_count = 0; ++ if ((total_compile_count++ % 10) == 0) { ++ std::cerr << "."; ++ } ++ } ++ compile_count_cond.notify_all(); ++} ++ ++static void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline) { ++ VK_LOG_DEBUG("ggml_pipeline_destroy_pipeline(" << pipeline->name << ")"); ++ for (auto& pool : pipeline->descriptor_pools) { ++ device.destroyDescriptorPool(pool); ++ } ++ pipeline->descriptor_pools.clear(); ++ pipeline->descriptor_sets.clear(); ++ pipeline->descriptor_set_idx = 0; ++ ++ device.destroyDescriptorSetLayout(pipeline->dsl); ++ ++ device.destroyPipelineLayout(pipeline->layout); ++ ++ device.destroyShaderModule(pipeline->shader_module); ++ ++ device.destroyPipeline(pipeline->pipeline); ++} ++ ++static void ggml_pipeline_request_descriptor_sets(vk_device& device, vk_pipeline& pipeline, uint32_t n) { ++ VK_LOG_DEBUG("ggml_pipeline_request_descriptor_sets(" << pipeline->name << ", " << n << ")"); ++ device->pipeline_descriptor_set_requirements[pipeline->name] += n; ++} ++ ++static void ggml_pipeline_allocate_descriptor_sets(vk_device& device) { ++ std::lock_guard guard(device->mutex); ++ ++ for (auto& pair : device->pipeline_descriptor_set_requirements) { ++ vk_pipeline pipeline = device->pipelines.at(pair.first).lock(); ++ const uint64_t n = pair.second; ++ ++ VK_LOG_DEBUG("ggml_pipeline_allocate_descriptor_sets(" << pipeline->name << ", " << n << ")"); ++ ++ if (pipeline->descriptor_sets.size() >= pipeline->descriptor_set_idx + n) { ++ // Enough descriptors are available ++ continue; ++ } ++ ++ uint32_t to_alloc = pipeline->descriptor_set_idx + n - pipeline->descriptor_sets.size(); ++ uint32_t pool_remaining = VK_DEVICE_DESCRIPTOR_POOL_SIZE - pipeline->descriptor_sets.size() % VK_DEVICE_DESCRIPTOR_POOL_SIZE; ++ uint32_t pool_idx = pipeline->descriptor_sets.size() / VK_DEVICE_DESCRIPTOR_POOL_SIZE; ++ ++ while (to_alloc > 0) { ++ const uint32_t alloc_count = std::min(pool_remaining, to_alloc); ++ to_alloc -= alloc_count; ++ pool_remaining = VK_DEVICE_DESCRIPTOR_POOL_SIZE; ++ ++ if (pool_idx >= pipeline->descriptor_pools.size()) { ++ vk::DescriptorPoolSize descriptor_pool_size(vk::DescriptorType::eStorageBuffer, pipeline->parameter_count * VK_DEVICE_DESCRIPTOR_POOL_SIZE); ++ vk::DescriptorPoolCreateInfo descriptor_pool_create_info({}, VK_DEVICE_DESCRIPTOR_POOL_SIZE, descriptor_pool_size); ++ pipeline->descriptor_pools.push_back(device->device.createDescriptorPool(descriptor_pool_create_info)); ++ } ++ ++ std::vector layouts(alloc_count); ++ for (uint32_t i = 0; i < alloc_count; i++) { ++ layouts[i] = pipeline->dsl; ++ } ++ vk::DescriptorSetAllocateInfo descriptor_set_alloc_info(pipeline->descriptor_pools[pool_idx], alloc_count, layouts.data()); ++ std::vector sets = device->device.allocateDescriptorSets(descriptor_set_alloc_info); ++ pipeline->descriptor_sets.insert(pipeline->descriptor_sets.end(), sets.begin(), sets.end()); ++ ++ pool_idx++; ++ } ++ } ++} ++ ++static void ggml_pipeline_cleanup(vk_pipeline& pipeline) { ++ VK_LOG_DEBUG("ggml_pipeline_cleanup(" << pipeline->name << ")"); ++ pipeline->descriptor_set_idx = 0; ++} ++ ++static vk::CommandBuffer ggml_vk_create_cmd_buffer(vk_device& device, vk_queue& q) { ++ VK_LOG_DEBUG("ggml_vk_create_cmd_buffer()"); ++ std::lock_guard guard(device->mutex); ++ ++ if (q.cmd_buffers.size() > q.cmd_buffer_idx) { ++ // Reuse command buffer ++ return q.cmd_buffers[q.cmd_buffer_idx++]; ++ } ++ ++ vk::CommandBufferAllocateInfo command_buffer_alloc_info( ++ q.pool, ++ vk::CommandBufferLevel::ePrimary, ++ 1); ++ const std::vector cmd_buffers = device->device.allocateCommandBuffers(command_buffer_alloc_info); ++ auto buf = cmd_buffers.front(); ++ ++ q.cmd_buffers.push_back(buf); ++ q.cmd_buffer_idx++; ++ ++ return buf; ++} ++ ++static vk_submission ggml_vk_create_submission(vk_device& device, vk_queue& q, std::vector wait_semaphores, std::vector signal_semaphores) { ++ VK_LOG_DEBUG("ggml_vk_create_submission()"); ++ vk_submission s; ++ s.buffer = ggml_vk_create_cmd_buffer(device, q); ++ s.wait_semaphores = std::move(wait_semaphores); ++ s.signal_semaphores = std::move(signal_semaphores); ++ return s; ++} ++ ++static void ggml_vk_submit(vk_context& ctx, vk::Fence fence) { ++ if (ctx->seqs.empty()) { ++ if (fence) { ++ ctx->q->queue.submit({}, fence); ++ } ++ return; ++ } ++ VK_LOG_DEBUG("ggml_vk_submit(" << ctx << ", " << fence << ")"); ++ ++ std::vector> tl_wait_vals; ++ std::vector> tl_signal_vals; ++ std::vector> tl_wait_semaphores; ++ std::vector> tl_signal_semaphores; ++ std::vector tl_submit_infos; ++ std::vector submit_infos; ++ int idx = -1; ++ std::vector> stage_flags; ++ ++ size_t reserve = 0; ++ ++ for (const auto& sequence : ctx->seqs) { ++ reserve += sequence.size(); ++ } ++ ++ // Pre-reserve vectors to prevent reallocation, which invalidates pointers ++ tl_wait_semaphores.reserve(reserve); ++ tl_wait_vals.reserve(reserve); ++ tl_signal_semaphores.reserve(reserve); ++ tl_signal_vals.reserve(reserve); ++ tl_submit_infos.reserve(reserve); ++ submit_infos.reserve(reserve); ++ stage_flags.reserve(reserve); ++ ++ for (const auto& sequence : ctx->seqs) { ++ for (const auto& submission : sequence) { ++ stage_flags.push_back({}); ++ idx++; ++ tl_wait_vals.push_back({}); ++ tl_wait_semaphores.push_back({}); ++ tl_signal_vals.push_back({}); ++ tl_signal_semaphores.push_back({}); ++ for (size_t i = 0; i < submission.wait_semaphores.size(); i++) { ++ stage_flags[idx].push_back(ctx->q->stage_flags); ++ tl_wait_vals[idx].push_back(submission.wait_semaphores[i].value); ++ tl_wait_semaphores[idx].push_back(submission.wait_semaphores[i].s); ++ } ++ for (size_t i = 0; i < submission.signal_semaphores.size(); i++) { ++ tl_signal_vals[idx].push_back(submission.signal_semaphores[i].value); ++ tl_signal_semaphores[idx].push_back(submission.signal_semaphores[i].s); ++ } ++ tl_submit_infos.push_back({ ++ (uint32_t) submission.wait_semaphores.size(), ++ tl_wait_vals[idx].data(), ++ (uint32_t) submission.signal_semaphores.size(), ++ tl_signal_vals[idx].data(), ++ }); ++ tl_submit_infos[idx].sType = vk::StructureType::eTimelineSemaphoreSubmitInfo; ++ tl_submit_infos[idx].pNext = nullptr; ++ vk::SubmitInfo si{ ++ (uint32_t) submission.wait_semaphores.size(), ++ tl_wait_semaphores[idx].data(), ++ stage_flags[idx].data(), ++ 1, ++ &submission.buffer, ++ (uint32_t) submission.signal_semaphores.size(), ++ tl_signal_semaphores[idx].data(), ++ }; ++ si.setPNext(&tl_submit_infos[idx]); ++ submit_infos.push_back(si); ++ } ++ } ++ ++ ctx->q->queue.submit(submit_infos, fence); ++ ++ ctx->seqs.clear(); ++} ++ ++static uint32_t ggml_vk_find_queue_family_index(std::vector& queue_family_props, const vk::QueueFlags& required, const vk::QueueFlags& avoid, int32_t compute_index, uint32_t min_num_queues) { ++ VK_LOG_DEBUG("ggml_vk_find_queue_family_index()"); ++ const uint32_t qfsize = queue_family_props.size(); ++ ++ // Try with avoid preferences first ++ for (uint32_t i = 0; i < qfsize; i++) { ++ if (queue_family_props[i].queueCount >= min_num_queues && (compute_index < 0 || i != (uint32_t) compute_index) && queue_family_props[i].queueFlags & required && !(queue_family_props[i].queueFlags & avoid)) { ++ return i; ++ } ++ } ++ ++ // Fall back to only required ++ for (size_t i = 0; i < qfsize; i++) { ++ if (queue_family_props[i].queueCount >= min_num_queues && (compute_index < 0 || i != (uint32_t) compute_index) && queue_family_props[i].queueFlags & required) { ++ return i; ++ } ++ } ++ ++ // Fall back to reusing compute queue ++ for (size_t i = 0; i < qfsize; i++) { ++ if (queue_family_props[i].queueCount >= min_num_queues && queue_family_props[i].queueFlags & required) { ++ return i; ++ } ++ } ++ ++ // Fall back to ignoring min_num_queries ++ for (size_t i = 0; i < qfsize; i++) { ++ if (queue_family_props[i].queueFlags & required) { ++ return i; ++ } ++ } ++ ++ // All commands that are allowed on a queue that supports transfer operations are also allowed on a queue that supports either graphics or compute operations. ++ // Thus, if the capabilities of a queue family include VK_QUEUE_GRAPHICS_BIT or VK_QUEUE_COMPUTE_BIT, then reporting the VK_QUEUE_TRANSFER_BIT capability separately for that queue family is optional. ++ if (compute_index >= 0) { ++ return compute_index; ++ } ++ ++ std::cerr << "ggml_vulkan: No suitable queue family index found." << std::endl; ++ ++ for(auto &q_family : queue_family_props) { ++ std::cerr << "Queue number: " + std::to_string(q_family.queueCount) << " flags: " + to_string(q_family.queueFlags) << std::endl; ++ } ++ abort(); ++} ++ ++static void ggml_vk_create_queue(vk_device& device, vk_queue& q, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) { ++ VK_LOG_DEBUG("ggml_vk_create_queue()"); ++ std::lock_guard guard(device->mutex); ++ ++ q.queue_family_index = queue_family_index; ++ q.transfer_only = transfer_only; ++ ++ vk::CommandPoolCreateInfo command_pool_create_info_compute(vk::CommandPoolCreateFlags(VK_COMMAND_POOL_CREATE_TRANSIENT_BIT), queue_family_index); ++ q.pool = device->device.createCommandPool(command_pool_create_info_compute); ++ ++ q.cmd_buffer_idx = 0; ++ ++ q.queue = device->device.getQueue(queue_family_index, queue_index); ++ ++ q.stage_flags = stage_flags; ++} ++ ++static vk_context ggml_vk_create_context(ggml_backend_vk_context * ctx, vk_queue& q) { ++ vk_context result = std::make_shared(); ++ VK_LOG_DEBUG("ggml_vk_create_context(" << result << ")"); ++ ctx->gc.contexts.emplace_back(result); ++ result->q = &q; ++ return result; ++} ++ ++static vk_context ggml_vk_create_temporary_context(vk_queue& q) { ++ vk_context result = std::make_shared(); ++ VK_LOG_DEBUG("ggml_vk_create_temporary_context(" << result << ")"); ++ result->q = &q; ++ return result; ++} ++ ++static vk_semaphore * ggml_vk_create_binary_semaphore(ggml_backend_vk_context * ctx) { ++ VK_LOG_DEBUG("ggml_vk_create_timeline_semaphore()"); ++ vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eBinary, 0 }; ++ vk::SemaphoreCreateInfo ci{}; ++ ci.setPNext(&tci); ++ vk::Semaphore semaphore = ctx->device->device.createSemaphore(ci); ++ ctx->gc.semaphores.push_back({ semaphore, 0 }); ++ return &ctx->gc.semaphores[ctx->gc.semaphores.size() - 1]; ++} ++ ++static vk_semaphore * ggml_vk_create_timeline_semaphore(ggml_backend_vk_context * ctx) { ++ VK_LOG_DEBUG("ggml_vk_create_timeline_semaphore()"); ++ if (ctx->semaphore_idx >= ctx->gc.tl_semaphores.size()) { ++ vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eTimeline, 0 }; ++ vk::SemaphoreCreateInfo ci{}; ++ ci.setPNext(&tci); ++ vk::Semaphore semaphore = ctx->device->device.createSemaphore(ci); ++ ctx->gc.tl_semaphores.push_back({ semaphore, 0 }); ++ } ++ return &ctx->gc.tl_semaphores[ctx->semaphore_idx++]; ++} ++ ++static vk::Event ggml_vk_create_event(ggml_backend_vk_context * ctx) { ++ if (ctx->event_idx >= ctx->gc.events.size()) { ++ ctx->gc.events.push_back(ctx->device->device.createEvent({})); ++ } ++ return ctx->gc.events[ctx->event_idx++]; ++} ++ ++static void ggml_vk_queue_cleanup(vk_device& device, vk_queue& q) { ++ VK_LOG_DEBUG("ggml_vk_queue_cleanup()"); ++ std::lock_guard guard(device->mutex); ++ ++ // Requires command buffers to be done ++ device->device.resetCommandPool(q.pool); ++ q.cmd_buffer_idx = 0; ++} ++ ++static uint32_t find_properties(const vk::PhysicalDeviceMemoryProperties* mem_props, vk::MemoryRequirements* mem_req, vk::MemoryPropertyFlags flags) { ++ for (uint32_t i = 0; i < mem_props->memoryTypeCount; ++i) { ++ vk::MemoryType memory_type = mem_props->memoryTypes[i]; ++ if ((mem_req->memoryTypeBits & ((uint64_t)1 << i)) && ++ (flags & memory_type.propertyFlags) == flags && ++ mem_props->memoryHeaps[memory_type.heapIndex].size >= mem_req->size) { ++ return static_cast(i); ++ } ++ } ++ return UINT32_MAX; ++} ++ ++static vk_buffer ggml_vk_create_buffer(vk_device& device, size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags fallback_flags = vk::MemoryPropertyFlags(0)) { ++ VK_LOG_DEBUG("ggml_vk_create_buffer(" << device->name << ", " << size << ", " << to_string(req_flags) << ", " << to_string(fallback_flags) << ")"); ++ if (size > device->max_memory_allocation_size) { ++ throw vk::OutOfDeviceMemoryError("Requested buffer size exceeds device memory allocation limit"); ++ } ++ ++ std::lock_guard guard(device->mutex); ++ ++ vk_buffer buf = std::make_shared(); ++ ++ if (size == 0) { ++ buf->size = 0; ++ return buf; ++ } ++ ++ vk::BufferCreateInfo buffer_create_info{ ++ vk::BufferCreateFlags(), ++ size, ++ vk::BufferUsageFlagBits::eStorageBuffer | vk::BufferUsageFlagBits::eTransferSrc | vk::BufferUsageFlagBits::eTransferDst, ++ vk::SharingMode::eExclusive, ++ 0, ++ nullptr, ++ }; ++ ++ buf->buffer = device->device.createBuffer(buffer_create_info); ++ ++ vk::MemoryRequirements mem_req = device->device.getBufferMemoryRequirements(buf->buffer); ++ ++ vk::PhysicalDeviceMemoryProperties mem_props = device->physical_device.getMemoryProperties(); ++ ++ uint32_t memory_type_index = UINT32_MAX; ++ ++ memory_type_index = find_properties(&mem_props, &mem_req, req_flags); ++ buf->memory_property_flags = req_flags; ++ ++ if (memory_type_index == UINT32_MAX && fallback_flags) { ++ memory_type_index = find_properties(&mem_props, &mem_req, fallback_flags); ++ buf->memory_property_flags = fallback_flags; ++ } ++ ++ if (memory_type_index == UINT32_MAX) { ++ device->device.destroyBuffer(buf->buffer); ++ throw vk::OutOfDeviceMemoryError("No suitable memory type found"); ++ } ++ ++ try { ++ buf->device_memory = device->device.allocateMemory({ mem_req.size, memory_type_index }); ++ } catch (const vk::SystemError& e) { ++ if (buf->memory_property_flags != fallback_flags) { ++ // Try again with fallback flags ++ memory_type_index = find_properties(&mem_props, &mem_req, fallback_flags); ++ buf->memory_property_flags = fallback_flags; ++ ++ try { ++ buf->device_memory = device->device.allocateMemory({ mem_req.size, memory_type_index }); ++ } ++ catch (const vk::SystemError& e) { ++ device->device.destroyBuffer(buf->buffer); ++ throw e; ++ } ++ } else { ++ // Out of Host/Device memory, clean up buffer ++ device->device.destroyBuffer(buf->buffer); ++ throw e; ++ } ++ } ++ buf->ptr = nullptr; ++ ++ if (buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { ++ buf->ptr = device->device.mapMemory(buf->device_memory, 0, VK_WHOLE_SIZE); ++ } ++ ++ device->device.bindBufferMemory(buf->buffer, buf->device_memory, 0); ++ ++ buf->device = device; ++ buf->size = size; ++ ++#ifdef GGML_VULKAN_MEMORY_DEBUG ++ device->memory_logger->log_allocation(buf, size); ++#endif ++ ++ return buf; ++} ++ ++static vk_buffer ggml_vk_create_buffer_check(vk_device& device, size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags fallback_flags = vk::MemoryPropertyFlags(0)) { ++ try { ++ return ggml_vk_create_buffer(device, size, req_flags, fallback_flags); ++ } catch (const vk::SystemError& e) { ++ std::cerr << "ggml_vulkan: Memory allocation of size " << size << " failed." << std::endl; ++ std::cerr << "ggml_vulkan: " << e.what() << std::endl; ++ throw e; ++ } ++} ++ ++static vk_buffer ggml_vk_create_buffer_device(vk_device& device, size_t size) { ++ vk_buffer buf; ++ try { ++ if (device->uma) { ++ // Fall back to host memory type ++ buf = ggml_vk_create_buffer(device, size, vk::MemoryPropertyFlagBits::eDeviceLocal, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); ++ } else { ++ // use rebar if available, otherwise fallback to device only visible memory ++ buf = ggml_vk_create_buffer(device, size, vk::MemoryPropertyFlagBits::eDeviceLocal | vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent, vk::MemoryPropertyFlagBits::eDeviceLocal); ++ } ++ } catch (const vk::SystemError& e) { ++ std::cerr << "ggml_vulkan: Device memory allocation of size " << size << " failed." << std::endl; ++ std::cerr << "ggml_vulkan: " << e.what() << std::endl; ++ throw e; ++ } ++ ++ return buf; ++} ++ ++static void ggml_vk_destroy_buffer(vk_buffer& buf) { ++ if (buf == nullptr) { ++ return; ++ } ++ ++#ifdef GGML_VULKAN_MEMORY_DEBUG ++ if (buf->device != nullptr) { ++ buf->device->memory_logger->log_deallocation(buf); ++ } ++#endif ++ ++ buf.reset(); ++} ++ ++static vk_subbuffer ggml_vk_subbuffer(vk_buffer& buf) { ++ return { buf, 0, VK_WHOLE_SIZE }; ++} ++ ++static void ggml_vk_sync_buffers(vk_context& ctx) { ++ VK_LOG_DEBUG("ggml_vk_sync_buffers()"); ++ ++ const bool transfer_queue = ctx->q->transfer_only; ++ ++ ctx->s->buffer.pipelineBarrier( ++ ctx->q->stage_flags, ++ ctx->q->stage_flags, ++ {}, ++ { { ++ { !transfer_queue ? (vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) : (vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) }, ++ { !transfer_queue ? (vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) : (vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) } ++ } }, ++ {}, ++ {} ++ ); ++} ++ ++static void ggml_vk_wait_events(vk_context& ctx, std::vector&& events) { ++ VK_LOG_DEBUG("ggml_vk_wait_events()"); ++ if (events.empty()) { ++ return; ++ } ++ ++ ctx->s->buffer.waitEvents( ++ events, ++ ctx->q->stage_flags, ++ ctx->q->stage_flags, ++ {}, ++ {}, ++ {} ++ ); ++} ++ ++// number of rows/cols for flash attention shader ++static constexpr uint32_t flash_attention_num_small_rows = 32; ++static std::array fa_rows_cols(uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) { ++ GGML_UNUSED(clamp); ++ ++ // small rows, large cols ++ if (small_rows) { ++ return {flash_attention_num_small_rows, 128}; ++ } ++ // small cols to reduce register count ++ if (ggml_is_quantized(type) || D == 256) { ++ return {64, 32}; ++ } ++ return {64, 64}; ++}; ++ ++static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vector& warptile, bool mul_mat_id) { ++ // Needs to be kept up to date on shader changes ++ const uint32_t bank_conflict_offset = device->coopmat_support ? 8 : 1; ++ const uint32_t type_size = device->fp16 ? sizeof(ggml_fp16_t) : sizeof(float); ++ const uint32_t warps = warptile[0] / warptile[10]; ++ ++ const uint32_t load_bufs = (warptile[1] + warptile[2]) * (warptile[3] + bank_conflict_offset) * type_size; ++ const uint32_t mmid_row_ids = mul_mat_id ? 3072 * sizeof(uint32_t) : 0; ++ const uint32_t coopmat_stage = device->coopmat_support ? warptile[7] * warptile[8] / warps * sizeof(float) : 0; ++ ++ return (load_bufs + mmid_row_ids + coopmat_stage) <= device->properties.limits.maxComputeSharedMemorySize; ++} ++ ++static void ggml_vk_load_shaders(vk_device& device) { ++ VK_LOG_DEBUG("ggml_vk_load_shaders(" << device->name << ")"); ++ ++ std::cerr << "ggml_vulkan: Compiling shaders"; ++ ++ // some shaders have a minimum subgroup size ++ const uint32_t subgroup_size_16 = std::max(device->subgroup_size, 16u); ++ const uint32_t subgroup_size_32 = std::max(device->subgroup_size, 32u); ++ ++ // mulmat ++ std::vector l_warptile, m_warptile, s_warptile, ++ l_warptile_mmq, m_warptile_mmq, s_warptile_mmq, ++ l_warptile_mmq_k, m_warptile_mmq_k, s_warptile_mmq_k, ++ l_warptile_mmqid, m_warptile_mmqid, s_warptile_mmqid; ++ std::array l_wg_denoms, m_wg_denoms, s_wg_denoms, ++ l_mmq_wg_denoms, m_mmq_wg_denoms, s_mmq_wg_denoms, ++ l_mmq_wg_denoms_k, m_mmq_wg_denoms_k, s_mmq_wg_denoms_k, ++ l_mmqid_wg_denoms, m_mmqid_wg_denoms, s_mmqid_wg_denoms; ++ ++ uint32_t l_align, m_align, s_align; ++ if (device->coopmat2) { ++ // spec constants and tile sizes for non-quant matmul/matmul_id ++ l_warptile = { 256, 128, 256, 64 }; ++ m_warptile = { 256, 128, 128, 64 }; ++ s_warptile = { 128, 64, 64, 64 }; ++ l_wg_denoms = {128, 256, 1 }; ++ m_wg_denoms = {128, 128, 1 }; ++ s_wg_denoms = { 64, 64, 1 }; ++ ++ // spec constants and tile sizes for quant matmul (non-Qi_K) ++ l_warptile_mmq = { 256, 128, 256, 64 }; ++ m_warptile_mmq = { 256, 128, 128, 64 }; ++ s_warptile_mmq = { 256, 128, 128, 64 }; ++ l_mmq_wg_denoms = { 128, 256, 1 }; ++ m_mmq_wg_denoms = { 128, 128, 1 }; ++ s_mmq_wg_denoms = { 128, 128, 1 }; ++ ++ // spec constants and tile sizes for quant matmul (Qi_K) ++ l_warptile_mmq_k = { 256, 128, 512, 16 }; ++ m_warptile_mmq_k = { 256, 128, 256, 16 }; ++ s_warptile_mmq_k = { 256, 32, 128, 64 }; ++ l_mmq_wg_denoms_k = { 128, 512, 1 }; ++ m_mmq_wg_denoms_k = { 128, 256, 1 }; ++ s_mmq_wg_denoms_k = { 32, 128, 1 }; ++ ++ // spec constants and tile sizes for quant matmul_id ++ l_warptile_mmqid = { 256, 128, 128, 16 }; ++ m_warptile_mmqid = { 256, 128, 64, 16 }; ++ s_warptile_mmqid = { 256, 64, 64, 16 }; ++ l_mmqid_wg_denoms = { 128, 128, 1 }; ++ m_mmqid_wg_denoms = { 128, 64, 1 }; ++ s_mmqid_wg_denoms = { 64, 64, 1 }; ++ ++ l_align = 128; ++ m_align = 64; ++ s_align = 32; ++ } else { ++ // Matrix cores require different warp group sizes ++ const uint32_t tm_l = device->coopmat_support ? device->coopmat_m : 4; ++ const uint32_t tm_m = device->coopmat_support ? device->coopmat_m : 4; ++ const uint32_t tm_s = device->coopmat_support ? device->coopmat_m : 2; ++ const uint32_t tn_l = device->coopmat_support ? device->coopmat_n : 4; ++ const uint32_t tn_m = device->coopmat_support ? device->coopmat_n : 2; ++ const uint32_t tn_s = device->coopmat_support ? device->coopmat_n : 2; ++ const uint32_t tk_l = device->coopmat_support ? device->coopmat_k : 1; ++ const uint32_t tk_m = device->coopmat_support ? device->coopmat_k : 1; ++ const uint32_t tk_s = device->coopmat_support ? device->coopmat_k : 1; ++ ++ l_warptile = { 128, 128, 128, 16, device->subgroup_size * 2, 64, 2, tm_l, tn_l, tk_l, device->subgroup_size }; ++ m_warptile = { 128, 64, 64, 16, device->subgroup_size, 32, 2, tm_m, tn_m, tk_m, device->subgroup_size }; ++ s_warptile = { subgroup_size_16, 32, 32, 16, 32, 32, 2, tm_s, tn_s, tk_s, device->subgroup_size }; ++ ++ l_warptile_mmq = { 128, 128, 128, 32, device->subgroup_size * 2, 64, 2, tm_l, tn_l, tk_l, device->subgroup_size }; ++ m_warptile_mmq = { 128, 64, 64, 32, device->subgroup_size, 32, 2, tm_m, tn_m, tk_m, device->subgroup_size }; ++ s_warptile_mmq = { subgroup_size_32, 32, 32, 32, 32, 32, 2, tm_s, tn_s, tk_s, device->subgroup_size }; ++ ++ l_mmq_wg_denoms = l_wg_denoms = {128, 128, 1 }; ++ m_mmq_wg_denoms = m_wg_denoms = { 64, 64, 1 }; ++ s_mmq_wg_denoms = s_wg_denoms = { 32, 32, 1 }; ++ l_align = 128; ++ m_align = 64; ++ s_align = 32; ++ ++ // Fallback to smaller sizes if there's not enough shared memory. Given the current shaders ++ // and tile sizes, this should handle 16KB, 32KB, and 48KB+. ++ // This logic doesn't explicitly account for the 12KB row_ids in the mul_mat_mat_id shaders. ++ // But the numbers happen to work out for 32KB shared memory size that when using the medium ++ // size there's enough room for everything, and we assert for this. ++ uint32_t shmem_needed = (l_warptile[1] + l_warptile[2]) * (l_warptile[3] + 1) * sizeof(float); ++ if (shmem_needed > device->properties.limits.maxComputeSharedMemorySize) { ++ l_warptile = m_warptile; ++ l_wg_denoms = m_wg_denoms; ++ shmem_needed = (l_warptile[1] + l_warptile[2]) * (l_warptile[3] + 1) * sizeof(float); ++ GGML_ASSERT(shmem_needed <= device->properties.limits.maxComputeSharedMemorySize); ++ } ++ if (device->properties.limits.maxComputeSharedMemorySize >= 32768) { ++ // assert mul_mat_mat_id shaders will fit. ++ GGML_ASSERT(shmem_needed + 3072*4 <= device->properties.limits.maxComputeSharedMemorySize); ++ } ++ ++ shmem_needed = (l_warptile_mmq[1] + l_warptile_mmq[2]) * (l_warptile_mmq[3] + 1) * sizeof(float); ++ if (shmem_needed > device->properties.limits.maxComputeSharedMemorySize) { ++ if (device->properties.limits.maxComputeSharedMemorySize == 32768) { ++ l_warptile_mmq = m_warptile_mmq; ++ l_mmq_wg_denoms = m_mmq_wg_denoms; ++ } else { ++ l_warptile_mmq = s_warptile_mmq; ++ l_mmq_wg_denoms = s_mmq_wg_denoms; ++ } ++ shmem_needed = (l_warptile_mmq[1] + l_warptile_mmq[2]) * (l_warptile_mmq[3] + 1) * sizeof(float); ++ GGML_ASSERT(shmem_needed <= device->properties.limits.maxComputeSharedMemorySize); ++ } ++ if (device->properties.limits.maxComputeSharedMemorySize >= 32768) { ++ // assert mul_mat_mat_id shaders will fit. ++ GGML_ASSERT(shmem_needed + 3072*4 <= device->properties.limits.maxComputeSharedMemorySize); ++ } ++ // Disable medium and large matrix multiplication if not enough shared memory is available ++ // Check mmq warptiles as the largest configuration ++ // Throw an error if not enough for any matrix multiplication is available ++ if (!ggml_vk_matmul_shmem_support(device, s_warptile_mmq, false)) { ++ std::cerr << "ggml_vulkan: Error: Shared memory size too small for matrix multiplication." << std::endl; ++ throw std::runtime_error("Shared memory size too small for matrix multiplication."); ++ } else if (!ggml_vk_matmul_shmem_support(device, m_warptile_mmq, false)) { ++ device->mul_mat_m = false; ++ device->mul_mat_l = false; ++ } else if (!ggml_vk_matmul_shmem_support(device, l_warptile_mmq, false)) { ++ device->mul_mat_l = false; ++ } ++ ++ // Disable mul_mat_id if not enough shared memory is available ++ if (!ggml_vk_matmul_shmem_support(device, s_warptile_mmq, true)) { ++ device->mul_mat_id_s = false; ++ device->mul_mat_id_m = false; ++ device->mul_mat_id_l = false; ++ } else if (!ggml_vk_matmul_shmem_support(device, m_warptile_mmq, true)) { ++ device->mul_mat_id_m = false; ++ device->mul_mat_id_l = false; ++ } else if (!ggml_vk_matmul_shmem_support(device, l_warptile_mmq, true)) { ++ device->mul_mat_id_l = false; ++ } ++ } ++ ++ device->pipeline_matmul_f32 = std::make_shared(); ++ device->pipeline_matmul_f32_f16 = std::make_shared(); ++ ++ device->pipeline_matmul_id_f32 = std::make_shared(); ++ ++ std::vector> compiles; ++ auto const &ggml_vk_create_pipeline = [&](vk_device& device, vk_pipeline& pipeline, const std::string &name, size_t spv_size, const void* spv_data, const std::string &entrypoint, ++ uint32_t parameter_count, uint32_t push_constant_size, std::array wg_denoms, const std::vector& specialization_constants, ++ uint32_t align, bool disable_robustness = false, bool require_full_subgroups = false, uint32_t required_subgroup_size = 0) { ++ { ++ // wait until fewer than N compiles are in progress ++ uint32_t N = std::max(1u, std::thread::hardware_concurrency()); ++ std::unique_lock guard(compile_count_mutex); ++ while (compile_count >= N) { ++ compile_count_cond.wait(guard); ++ } ++ compile_count++; ++ } ++ compiles.push_back(std::async(ggml_vk_create_pipeline_func, std::ref(device), std::ref(pipeline), name, spv_size, spv_data, entrypoint, ++ parameter_count, push_constant_size, wg_denoms, specialization_constants, align, disable_robustness, require_full_subgroups, required_subgroup_size)); ++ }; ++ ++#if defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) ++ if (device->coopmat2) { ++ ++ auto const &fa_wg_denoms = [&](uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) -> std::array { ++ return {fa_rows_cols(D, clamp, type, small_rows)[0], 1, 1}; ++ }; ++ ++ auto const &fa_spec_constants = [&](uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) -> std::vector { ++ // For large number of rows, 128 invocations seems to work best. ++ // For small number of rows (e.g. N==1), 256 works better. But matrix granularity for 256 is 32, so we ++ // can't use 256 for D==80. ++ uint32_t wg_size = (small_rows && (D % 32) == 0) ? 256 : 128; ++ auto rows_cols = fa_rows_cols(D, clamp, type, small_rows); ++ return {wg_size, rows_cols[0], rows_cols[1], (D), clamp}; ++ }; ++ ++#define CREATE_FA2(TYPE, NAMELC, D) \ ++ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][0][0], "flash_attn_f32_f16_D" #D "_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,false), fa_spec_constants(D,1,TYPE,false), 1); \ ++ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][0][1], "flash_attn_f32_f16_D" #D "_aligned_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,false), fa_spec_constants(D,0,TYPE,false), fa_rows_cols(D,0,TYPE,false)[1]); \ ++ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][0][0], "flash_attn_f32_f16_D" #D "_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,false), fa_spec_constants(D,1,TYPE,false), 1); \ ++ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][0][1], "flash_attn_f32_f16_D" #D "_aligned_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,false), fa_spec_constants(D,0,TYPE,false), fa_rows_cols(D,0,TYPE,false)[1]); \ ++ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][1][0], "flash_attn_f32_f16_D" #D "_f16acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,true), fa_spec_constants(D,1,TYPE,true), 1); \ ++ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][1][1], "flash_attn_f32_f16_D" #D "_aligned_f16acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,true), fa_spec_constants(D,0,TYPE,true), fa_rows_cols(D,0,TYPE,true)[1]); \ ++ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][1][0], "flash_attn_f32_f16_D" #D "_f32acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,true), fa_spec_constants(D,1,TYPE,true), 1); \ ++ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][1][1], "flash_attn_f32_f16_D" #D "_aligned_f32acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,true), fa_spec_constants(D,0,TYPE,true), fa_rows_cols(D,0,TYPE,true)[1]); \ ++ ++#define CREATE_FA(TYPE, NAMELC) \ ++ CREATE_FA2(TYPE, NAMELC, 64) \ ++ CREATE_FA2(TYPE, NAMELC, 80) \ ++ CREATE_FA2(TYPE, NAMELC, 96) \ ++ CREATE_FA2(TYPE, NAMELC, 112) \ ++ CREATE_FA2(TYPE, NAMELC, 128) \ ++ CREATE_FA2(TYPE, NAMELC, 256) ++ ++ CREATE_FA(GGML_TYPE_F16, f16) ++ CREATE_FA(GGML_TYPE_Q4_0, q4_0) ++ CREATE_FA(GGML_TYPE_Q4_1, q4_1) ++ CREATE_FA(GGML_TYPE_Q5_0, q5_0) ++ CREATE_FA(GGML_TYPE_Q5_1, q5_1) ++ CREATE_FA(GGML_TYPE_Q8_0, q8_0) ++ // K dequants currently disabled because D dimension is rounded up to 256 and runs inefficiently ++ //CREATE_FA(GGML_TYPE_Q2_K, q2_k) ++ //CREATE_FA(GGML_TYPE_Q3_K, q3_k) ++ //CREATE_FA(GGML_TYPE_Q4_K, q4_k) ++ //CREATE_FA(GGML_TYPE_Q5_K, q5_k) ++ //CREATE_FA(GGML_TYPE_Q6_K, q6_k) ++ CREATE_FA(GGML_TYPE_IQ4_NL, iq4_nl) ++#undef CREATE_FA ++ ++ // Create 6 variants, {s,m,l}x{unaligned,aligned} ++#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ ++ ++ // Create 2 variants, {f16,f32} accumulator ++#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ ++ CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ ++ CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ ++ ++ CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3) ++ ++ CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3) ++ CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) ++ ++ CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) ++#undef CREATE_MM ++#undef CREATE_MM2 ++ } else ++#endif // defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) ++ if (device->coopmat_support) { ++ // Create 6 variants, {s,m,l}x{unaligned,aligned} ++#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ ++ if (device->mul_mat ## ID ## _l) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1, false, true); \ ++ if (device->mul_mat ## ID ## _m) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1, false, true); \ ++ if (device->mul_mat ## ID ## _s) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1, false, true); \ ++ if (device->mul_mat ## ID ## _l) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align, false, true); \ ++ if (device->mul_mat ## ID ## _m) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align, false, true); \ ++ if (device->mul_mat ## ID ## _s) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align, false, true); \ ++ ++ // Create 2 variants, {f16,f32} accumulator ++#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ ++ if (device->coopmat_acc_f16_support) { \ ++ CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ ++ } \ ++ if (device->coopmat_acc_f32_support) { \ ++ CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ ++ } \ ++ ++ CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ ++ if (device->coopmat_acc_f16_support) { ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ } else { ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ } ++ ++ // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. ++ if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { ++ CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); ++ CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); ++ CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); ++ ++ if (device->coopmat_acc_f16_support) { ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ } else { ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ } ++ } ++#undef CREATE_MM2 ++#undef CREATE_MM ++ } else if (device->fp16) { ++ // Create 6 variants, {s,m,l}x{unaligned,aligned} ++#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ ++ if (device->mul_mat ## ID ## _l) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ ++ if (device->mul_mat ## ID ## _m) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ ++ if (device->mul_mat ## ID ## _s) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ ++ if (device->mul_mat ## ID ## _l) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ ++ if (device->mul_mat ## ID ## _m) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ ++ if (device->mul_mat ## ID ## _s) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ ++ ++ // Create 2 variants, {f16,f32} accumulator ++#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ ++ CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ ++ CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ ++ ++ CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ ++ // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. ++ if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { ++ CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); ++ CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); ++ CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ } ++#undef CREATE_MM2 ++#undef CREATE_MM ++ } else { ++ // Create 6 variants, {s,m,l}x{unaligned,aligned} ++#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ ++ if (device->mul_mat ## ID ## _l) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ ++ if (device->mul_mat ## ID ## _m) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ ++ if (device->mul_mat ## ID ## _s) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ ++ if (device->mul_mat ## ID ## _l) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ ++ if (device->mul_mat ## ID ## _m) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ ++ if (device->mul_mat ## ID ## _s) \ ++ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ ++ ++ CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_matmul_f16.f32acc, matmul_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_matmul_f16_f32.f32acc, matmul_f16_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f32acc, matmul_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f32acc, matmul_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f32acc, matmul_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f32acc, matmul_q5_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f32acc, matmul_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f32acc, matmul_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f32acc, matmul_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f32acc, matmul_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f32acc, matmul_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f32acc, matmul_q6_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f32acc, matmul_iq4_nl_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); ++ ++ // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. ++ if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { ++ CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); ++ CREATE_MM(pipeline_matmul_id_f16.f32acc, matmul_id_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); ++ CREATE_MM(pipeline_matmul_id_f16_f32.f32acc, matmul_id_f16_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f32acc, matmul_id_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f32acc, matmul_id_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f32acc, matmul_id_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f32acc, matmul_id_q5_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f32acc, matmul_id_q6_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f32acc, matmul_id_iq4_nl_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); ++ } ++#undef CREATE_MM ++ } ++ ++ // mul mat vec ++ ++ // the number of rows computed per shader depends on GPU model and quant ++ uint32_t rm_stdq = 1; ++ uint32_t rm_kq = 2; ++ if (device->vendor_id == VK_VENDOR_ID_AMD) { ++ if (device->subgroup_min_size == 64 && device->subgroup_max_size == 64) { // GCN ++ rm_stdq = 2; ++ rm_kq = 4; ++ } ++ } else if (device->vendor_id == VK_VENDOR_ID_INTEL) ++ rm_stdq = 2; ++ ++ for (uint32_t i = 0; i < mul_mat_vec_max_cols; ++i) { ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_F32 ][i], "mul_mat_vec_f32_f32_f32_"+std::to_string(i+1), mul_mat_vec_f32_f32_f32_len, mul_mat_vec_f32_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f32_f32_"+std::to_string(i+1), mul_mat_vec_f16_f32_f32_len, mul_mat_vec_f16_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_0_f32_f32_len, mul_mat_vec_q4_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_1_f32_f32_len, mul_mat_vec_q4_1_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_0_f32_f32_len, mul_mat_vec_q5_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_1_f32_f32_len, mul_mat_vec_q5_1_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q8_0_f32_f32_len, mul_mat_vec_q8_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q2_k_f32_f32_len, mul_mat_vec_q2_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q3_k_f32_f32_len, mul_mat_vec_q3_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_k_f32_f32_len, mul_mat_vec_q4_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_k_f32_f32_len, mul_mat_vec_q5_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q6_k_f32_f32_len, mul_mat_vec_q6_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f32_f32_"+std::to_string(i+1), mul_mat_vec_iq4_nl_f32_f32_len, mul_mat_vec_iq4_nl_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq, i+1}, 1, true); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_F32 ][i], "mul_mat_vec_f32_f16_f32_"+std::to_string(i+1), mul_mat_vec_f32_f16_f32_len, mul_mat_vec_f32_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f16_f32_"+std::to_string(i+1), mul_mat_vec_f16_f16_f32_len, mul_mat_vec_f16_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_0_f16_f32_len, mul_mat_vec_q4_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_1_f16_f32_len, mul_mat_vec_q4_1_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_0_f16_f32_len, mul_mat_vec_q5_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_1_f16_f32_len, mul_mat_vec_q5_1_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q8_0_f16_f32_len, mul_mat_vec_q8_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q2_k_f16_f32_len, mul_mat_vec_q2_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q3_k_f16_f32_len, mul_mat_vec_q3_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_k_f16_f32_len, mul_mat_vec_q4_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_k_f16_f32_len, mul_mat_vec_q5_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q6_k_f16_f32_len, mul_mat_vec_q6_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f16_f32_"+std::to_string(i+1), mul_mat_vec_iq4_nl_f16_f32_len, mul_mat_vec_iq4_nl_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq, i+1}, 1, true); ++ } ++ ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_F32 ], "mul_mat_vec_id_f32_f32", mul_mat_vec_id_f32_f32_len, mul_mat_vec_id_f32_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {device->subgroup_size, 2}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_F16 ], "mul_mat_vec_id_f16_f32", mul_mat_vec_id_f16_f32_len, mul_mat_vec_id_f16_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {device->subgroup_size, 2}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_0], "mul_mat_vec_id_q4_0_f32", mul_mat_vec_id_q4_0_f32_len, mul_mat_vec_id_q4_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_1], "mul_mat_vec_id_q4_1_f32", mul_mat_vec_id_q4_1_f32_len, mul_mat_vec_id_q4_1_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_0], "mul_mat_vec_id_q5_0_f32", mul_mat_vec_id_q5_0_f32_len, mul_mat_vec_id_q5_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_1], "mul_mat_vec_id_q5_1_f32", mul_mat_vec_id_q5_1_f32_len, mul_mat_vec_id_q5_1_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q8_0], "mul_mat_vec_id_q8_0_f32", mul_mat_vec_id_q8_0_f32_len, mul_mat_vec_id_q8_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q2_K], "mul_mat_vec_id_q2_k_f32", mul_mat_vec_id_q2_k_f32_len, mul_mat_vec_id_q2_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q3_K], "mul_mat_vec_id_q3_k_f32", mul_mat_vec_id_q3_k_f32_len, mul_mat_vec_id_q3_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_K], "mul_mat_vec_id_q4_k_f32", mul_mat_vec_id_q4_k_f32_len, mul_mat_vec_id_q4_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_K], "mul_mat_vec_id_q5_k_f32", mul_mat_vec_id_q5_k_f32_len, mul_mat_vec_id_q5_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q6_K], "mul_mat_vec_id_q6_k_f32", mul_mat_vec_id_q6_k_f32_len, mul_mat_vec_id_q6_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_IQ4_NL], "mul_mat_vec_id_iq4_nl_f32", mul_mat_vec_id_iq4_nl_f32_len, mul_mat_vec_id_iq4_nl_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq}, 1, true); ++ ++ // dequant shaders ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_F32 ], "f32_to_f16", dequant_f32_len, dequant_f32_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_0], "dequant_q4_0", dequant_q4_0_len, dequant_q4_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_1], "dequant_q4_1", dequant_q4_1_len, dequant_q4_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_1], "dequant_q5_1", dequant_q5_1_len, dequant_q5_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_K], "dequant_q4_k", dequant_q4_k_len, dequant_q4_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_K], "dequant_q5_k", dequant_q5_k_len, dequant_q5_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q6_K], "dequant_q6_k", dequant_q6_k_len, dequant_q6_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ4_NL], "dequant_iq4_nl", dequant_iq4_nl_len, dequant_iq4_nl_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ++ ++ // get_rows ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_F32 ], "get_rows_f32", get_rows_f32_len, get_rows_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_F16 ], "get_rows_f16", get_rows_f16_len, get_rows_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_0], "get_rows_q4_0", get_rows_q4_0_len, get_rows_q4_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_1], "get_rows_q4_1", get_rows_q4_1_len, get_rows_q4_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_0], "get_rows_q5_0", get_rows_q5_0_len, get_rows_q5_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_1], "get_rows_q5_1", get_rows_q5_1_len, get_rows_q5_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q8_0], "get_rows_q8_0", get_rows_q8_0_len, get_rows_q8_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_IQ4_NL], "get_rows_iq4_nl", get_rows_iq4_nl_len, get_rows_iq4_nl_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_F32 ], "get_rows_f32_f32", get_rows_f32_f32_len, get_rows_f32_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_F16 ], "get_rows_f16_f32", get_rows_f16_f32_len, get_rows_f16_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_0], "get_rows_q4_0_f32", get_rows_q4_0_f32_len, get_rows_q4_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_1], "get_rows_q4_1_f32", get_rows_q4_1_f32_len, get_rows_q4_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_0], "get_rows_q5_0_f32", get_rows_q5_0_f32_len, get_rows_q5_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_1], "get_rows_q5_1_f32", get_rows_q5_1_f32_len, get_rows_q5_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q8_0], "get_rows_q8_0_f32", get_rows_q8_0_f32_len, get_rows_q8_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_IQ4_NL], "get_rows_iq4_nl_f32", get_rows_iq4_nl_f32_len, get_rows_iq4_nl_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_matmul_split_k_reduce, "split_k_reduce", split_k_reduce_len, split_k_reduce_data, "main", 2, 2 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_mul_mat_vec_p021_f16_f32, "mul_mat_vec_p021_f16_f32", mul_mat_vec_p021_f16_f32_len, mul_mat_vec_p021_f16_f32_data, "main", 3, 6 * sizeof(uint32_t), {1, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_mul_mat_vec_nc_f16_f32, "mul_mat_vec_nc_f16_f32", mul_mat_vec_nc_f16_f32_len, mul_mat_vec_nc_f16_f32_data, "main", 3, 7 * sizeof(uint32_t), {1, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_norm_f32, "norm_f32", norm_f32_len, norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_group_norm_f32, "group_norm_f32", group_norm_f32_len, group_norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_rms_norm_f32, "rms_norm_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f32, "cpy_f32_f32", cpy_f32_f32_len, cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f16, "cpy_f32_f16", cpy_f32_f16_len, cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_cpy_f16_f16, "cpy_f16_f16", cpy_f16_f16_len, cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f32, "contig_cpy_f32_f32", contig_cpy_f32_f32_len, contig_cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f16, "contig_cpy_f32_f16", contig_cpy_f32_f16_len, contig_cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f16_f16, "contig_cpy_f16_f16", contig_cpy_f16_f16_len, contig_cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_add_f32, "add_f32", add_f32_len, add_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_add_f32_norepeat, "add_f32_norepeat", add_f32_len, add_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_add_f16_f32_f16, "add_f16_f32_f16", add_f16_f32_f16_len, add_f16_f32_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_add_f16_f32_f16_norepeat, "add_f16_f32_f16_norepeat", add_f16_f32_f16_len, add_f16_f32_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_acc_f32, "acc_f32", acc_f32_len, acc_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_mul_f32, "mul_f32", mul_f32_len, mul_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_mul_f32_norepeat, "mul_f32_norepeat", mul_f32_len, mul_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_div_f32, "div_f32", div_f32_len, div_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_div_f32_norepeat, "div_f32_norepeat", div_f32_len, div_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_concat_f32, "concat_f32", concat_f32_len, concat_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_concat_f16, "concat_f16", concat_f16_len, concat_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_concat_i32, "concat_i32", concat_i32_len, concat_i32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_upscale_f32, "upscale_f32", upscale_f32_len, upscale_f32_data, "main", 2, sizeof(vk_op_upscale_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_scale_f32, "scale_f32", scale_f32_len, scale_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_sqr_f32, "sqr_f32", sqr_f32_len, sqr_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_sin_f32, "sin_f32", sin_f32_len, sin_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_cos_f32, "cos_f32", cos_f32_len, cos_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_clamp_f32, "clamp_f32", clamp_f32_len, clamp_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_pad_f32, "pad_f32", pad_f32_len, pad_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_repeat_f32, "repeat_f32", repeat_f32_len, repeat_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_gelu_f32, "gelu_f32", gelu_f32_len, gelu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_gelu_quick_f32, "gelu_quick_f32", gelu_quick_f32_len, gelu_quick_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_silu_f32, "silu_f32", silu_f32_len, silu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_relu_f32, "relu_f32", relu_f32_len, relu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_leaky_relu_f32, "leaky_relu_f32", leaky_relu_f32_len, leaky_relu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_tanh_f32, "tanh_f32", tanh_f32_len, tanh_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_diag_mask_inf_f32, "diag_mask_inf_f32", diag_mask_inf_f32_len, diag_mask_inf_f32_data, "main", 2, sizeof(vk_op_diag_mask_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32, "soft_max_f32", soft_max_f32_len, soft_max_f32_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_wg512, "soft_max_f32_wg512", soft_max_f32_len, soft_max_f32_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 512 }, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_f16, "soft_max_f32_f16", soft_max_f32_f16_len, soft_max_f32_f16_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_f16_wg512, "soft_max_f32_f16_wg512", soft_max_f32_f16_len, soft_max_f32_f16_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 512 }, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f32, "rope_norm_f32", rope_norm_f32_len, rope_norm_f32_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f32, "rope_neox_f32", rope_neox_f32_len, rope_neox_f32_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); ++ ++ if (device->float_controls_rte_fp16) { ++ ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f16, "rope_norm_f16", rope_norm_f16_rte_len, rope_norm_f16_rte_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f16, "rope_neox_f16", rope_neox_f16_rte_len, rope_neox_f16_rte_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); ++ } else { ++ ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f16, "rope_norm_f16", rope_norm_f16_len, rope_norm_f16_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); ++ ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f16, "rope_neox_f16", rope_neox_f16_len, rope_neox_f16_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); ++ } ++ ++ ggml_vk_create_pipeline(device, device->pipeline_argsort_f32, "argsort_f32", argsort_f32_len, argsort_f32_data, "main", 2, sizeof(vk_op_argsort_push_constants), {1024, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_im2col_f32, "im2col_f32", im2col_f32_len, im2col_f32_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); ++ if (device->float_controls_rte_fp16) { ++ ggml_vk_create_pipeline(device, device->pipeline_im2col_f32_f16, "im2col_f32_f16", im2col_f32_f16_rte_len, im2col_f32_f16_rte_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); ++ } else { ++ ggml_vk_create_pipeline(device, device->pipeline_im2col_f32_f16, "im2col_f32_f16", im2col_f32_f16_len, im2col_f32_f16_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); ++ } ++ ++ ggml_vk_create_pipeline(device, device->pipeline_timestep_embedding_f32, "timestep_embedding_f32", timestep_embedding_f32_len, timestep_embedding_f32_data, "main", 2, sizeof(vk_op_timestep_embedding_push_constants), {256, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_pool2d_f32, "pool2d_f32", pool2d_f32_len, pool2d_f32_data, "main", 2, sizeof(vk_op_pool2d_push_constants), {512, 1, 1}, {}, 1); ++ ++ ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv6_f32, "rwkv_wkv6_f32", rwkv_wkv6_f32_len, rwkv_wkv6_f32_data, "main", 7, sizeof(vk_op_rwkv_wkv6_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); ++ ++ for (auto &c : compiles) { ++ c.wait(); ++ } ++ std::cerr << "Done!" << std::endl; ++} ++ ++static bool ggml_vk_khr_cooperative_matrix_support(const vk::PhysicalDeviceProperties& props, const vk::PhysicalDeviceDriverProperties& driver_props); ++ ++static vk_device ggml_vk_get_device(size_t idx) { ++ VK_LOG_DEBUG("ggml_vk_get_device(" << idx << ")"); ++ ++ if (vk_instance.devices[idx] == nullptr) { ++ VK_LOG_DEBUG("Initializing new vk_device"); ++ vk_device device = std::make_shared(); ++ vk_instance.devices[idx] = device; ++ ++#ifdef GGML_VULKAN_MEMORY_DEBUG ++ device->memory_logger = std::unique_ptr(new vk_memory_logger()); ++#endif ++#ifdef GGML_VULKAN_PERF ++ device->perf_logger = std::unique_ptr(new vk_perf_logger()); ++#endif ++ ++ size_t dev_num = vk_instance.device_indices[idx]; ++ ++ std::vector physical_devices = vk_instance.instance.enumeratePhysicalDevices(); ++ ++ if (dev_num >= physical_devices.size()) { ++ std::cerr << "ggml_vulkan: Device with index " << dev_num << " does not exist." << std::endl; ++ throw std::runtime_error("Device not found"); ++ } ++ ++ device->physical_device = physical_devices[dev_num]; ++ const std::vector ext_props = device->physical_device.enumerateDeviceExtensionProperties(); ++ ++ bool fp16_storage = false; ++ bool fp16_compute = false; ++ bool maintenance4_support = false; ++ bool sm_builtins = false; ++ bool amd_shader_core_properties2 = false; ++ bool pipeline_robustness = false; ++ bool coopmat2_support = false; ++ device->coopmat_support = false; ++ ++ // Check if maintenance4 is supported ++ for (const auto& properties : ext_props) { ++ if (strcmp("VK_KHR_maintenance4", properties.extensionName) == 0) { ++ maintenance4_support = true; ++ } else if (strcmp("VK_KHR_16bit_storage", properties.extensionName) == 0) { ++ fp16_storage = true; ++ } else if (strcmp("VK_KHR_shader_float16_int8", properties.extensionName) == 0) { ++ fp16_compute = true; ++ } else if (strcmp("VK_NV_shader_sm_builtins", properties.extensionName) == 0) { ++ sm_builtins = true; ++ } else if (strcmp("VK_AMD_shader_core_properties2", properties.extensionName) == 0) { ++ amd_shader_core_properties2 = true; ++ } else if (strcmp("VK_EXT_pipeline_robustness", properties.extensionName) == 0) { ++ pipeline_robustness = true; ++ } else if (strcmp("VK_EXT_subgroup_size_control", properties.extensionName) == 0) { ++ device->subgroup_size_control = true; ++ } else if (strcmp("VK_KHR_cooperative_matrix", properties.extensionName) == 0 && ++ !getenv("GGML_VK_DISABLE_COOPMAT")) { ++ device->coopmat_support = true; ++ device->coopmat_m = 0; ++ device->coopmat_n = 0; ++ device->coopmat_k = 0; ++ } else if (strcmp("VK_NV_cooperative_matrix2", properties.extensionName) == 0 && ++ !getenv("GGML_VK_DISABLE_COOPMAT2")) { ++ coopmat2_support = true; ++ } ++ } ++ ++ vk::PhysicalDeviceProperties2 props2; ++ vk::PhysicalDeviceMaintenance3Properties props3; ++ vk::PhysicalDeviceMaintenance4Properties props4; ++ vk::PhysicalDeviceSubgroupProperties subgroup_props; ++ vk::PhysicalDeviceDriverProperties driver_props; ++ vk::PhysicalDeviceShaderSMBuiltinsPropertiesNV sm_props; ++ vk::PhysicalDeviceShaderCoreProperties2AMD amd_shader_core_properties2_props; ++ vk::PhysicalDeviceVulkan12Properties vk12_props; ++ vk::PhysicalDeviceSubgroupSizeControlPropertiesEXT subgroup_size_control_props; ++ ++ props2.pNext = &props3; ++ props3.pNext = &subgroup_props; ++ subgroup_props.pNext = &driver_props; ++ driver_props.pNext = &vk12_props; ++ ++ VkBaseOutStructure * last_struct = (VkBaseOutStructure *)&vk12_props; ++ ++ if (maintenance4_support) { ++ last_struct->pNext = (VkBaseOutStructure *)&props4; ++ last_struct = (VkBaseOutStructure *)&props4; ++ } ++ if (sm_builtins) { ++ last_struct->pNext = (VkBaseOutStructure *)&sm_props; ++ last_struct = (VkBaseOutStructure *)&sm_props; ++ } ++ if (amd_shader_core_properties2) { ++ last_struct->pNext = (VkBaseOutStructure *)&amd_shader_core_properties2_props; ++ last_struct = (VkBaseOutStructure *)&amd_shader_core_properties2_props; ++ } ++ if (device->subgroup_size_control) { ++ last_struct->pNext = (VkBaseOutStructure *)&subgroup_size_control_props; ++ last_struct = (VkBaseOutStructure *)&subgroup_size_control_props; ++ } ++ ++#if defined(VK_NV_cooperative_matrix2) ++ vk::PhysicalDeviceCooperativeMatrix2PropertiesNV coopmat2_props; ++ if (coopmat2_support) { ++ last_struct->pNext = (VkBaseOutStructure *)&coopmat2_props; ++ last_struct = (VkBaseOutStructure *)&coopmat2_props; ++ } ++#endif ++ ++ device->physical_device.getProperties2(&props2); ++ device->properties = props2.properties; ++ ++ const char* GGML_VK_FORCE_MAX_ALLOCATION_SIZE = getenv("GGML_VK_FORCE_MAX_ALLOCATION_SIZE"); ++ ++ if (GGML_VK_FORCE_MAX_ALLOCATION_SIZE != nullptr) { ++ device->max_memory_allocation_size = std::stoul(GGML_VK_FORCE_MAX_ALLOCATION_SIZE); ++ } else if (maintenance4_support) { ++ device->max_memory_allocation_size = std::min(props3.maxMemoryAllocationSize, props4.maxBufferSize); ++ } else { ++ device->max_memory_allocation_size = props3.maxMemoryAllocationSize; ++ } ++ ++ device->vendor_id = device->properties.vendorID; ++ device->subgroup_size = subgroup_props.subgroupSize; ++ device->uma = device->properties.deviceType == vk::PhysicalDeviceType::eIntegratedGpu; ++ if (sm_builtins) { ++ device->shader_core_count = sm_props.shaderSMCount; ++ } else if (amd_shader_core_properties2) { ++ device->shader_core_count = amd_shader_core_properties2_props.activeComputeUnitCount; ++ } else { ++ device->shader_core_count = 0; ++ } ++ device->float_controls_rte_fp16 = vk12_props.shaderRoundingModeRTEFloat16; ++ ++ const bool force_disable_f16 = getenv("GGML_VK_DISABLE_F16") != nullptr; ++ ++ device->fp16 = !force_disable_f16 && fp16_storage && fp16_compute; ++ ++ if (!ggml_vk_khr_cooperative_matrix_support(device->properties, driver_props)) { ++ device->coopmat_support = false; ++ } ++ ++ std::vector queue_family_props = device->physical_device.getQueueFamilyProperties(); ++ ++ // Try to find a non-graphics compute queue and transfer-focused queues ++ const uint32_t compute_queue_family_index = ggml_vk_find_queue_family_index(queue_family_props, vk::QueueFlagBits::eCompute, vk::QueueFlagBits::eGraphics, -1, 1); ++ const uint32_t transfer_queue_family_index = ggml_vk_find_queue_family_index(queue_family_props, vk::QueueFlagBits::eTransfer, vk::QueueFlagBits::eCompute | vk::QueueFlagBits::eGraphics, compute_queue_family_index, 1); ++ ++ const float priorities[] = { 1.0f, 1.0f }; ++ device->single_queue = compute_queue_family_index == transfer_queue_family_index && queue_family_props[compute_queue_family_index].queueCount == 1; ++ ++ std::vector device_queue_create_infos; ++ if (compute_queue_family_index != transfer_queue_family_index) { ++ device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities}); ++ device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), transfer_queue_family_index, 1, priorities + 1}); ++ } else if(!device->single_queue) { ++ device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 2, priorities}); ++ } else { ++ device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities}); ++ } ++ vk::DeviceCreateInfo device_create_info; ++ std::vector device_extensions; ++ vk::PhysicalDeviceFeatures device_features = device->physical_device.getFeatures(); ++ ++ VkPhysicalDeviceFeatures2 device_features2; ++ device_features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2; ++ device_features2.pNext = nullptr; ++ device_features2.features = (VkPhysicalDeviceFeatures)device_features; ++ ++ VkPhysicalDeviceVulkan11Features vk11_features; ++ vk11_features.pNext = nullptr; ++ vk11_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES; ++ device_features2.pNext = &vk11_features; ++ ++ VkPhysicalDeviceVulkan12Features vk12_features; ++ vk12_features.pNext = nullptr; ++ vk12_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES; ++ vk11_features.pNext = &vk12_features; ++ ++ last_struct = (VkBaseOutStructure *)&vk12_features; ++ ++ VkPhysicalDevicePipelineRobustnessFeaturesEXT pl_robustness_features; ++ pl_robustness_features.pNext = nullptr; ++ pl_robustness_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PIPELINE_ROBUSTNESS_FEATURES_EXT; ++ pl_robustness_features.pipelineRobustness = VK_FALSE; ++ ++ if (pipeline_robustness) { ++ last_struct->pNext = (VkBaseOutStructure *)&pl_robustness_features; ++ last_struct = (VkBaseOutStructure *)&pl_robustness_features; ++ device_extensions.push_back("VK_EXT_pipeline_robustness"); ++ } ++ ++ VkPhysicalDeviceSubgroupSizeControlFeaturesEXT subgroup_size_control_features; ++ subgroup_size_control_features.pNext = nullptr; ++ subgroup_size_control_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES_EXT; ++ subgroup_size_control_features.computeFullSubgroups = false; ++ subgroup_size_control_features.subgroupSizeControl = false; ++ ++ if (device->subgroup_size_control) { ++ last_struct->pNext = (VkBaseOutStructure *)&subgroup_size_control_features; ++ last_struct = (VkBaseOutStructure *)&subgroup_size_control_features; ++ } ++ ++ VkPhysicalDeviceCooperativeMatrixFeaturesKHR coopmat_features; ++ coopmat_features.pNext = nullptr; ++ coopmat_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR; ++ coopmat_features.cooperativeMatrix = VK_FALSE; ++ ++ if (device->coopmat_support) { ++ last_struct->pNext = (VkBaseOutStructure *)&coopmat_features; ++ last_struct = (VkBaseOutStructure *)&coopmat_features; ++ } ++ ++#if defined(VK_NV_cooperative_matrix2) ++ VkPhysicalDeviceCooperativeMatrix2FeaturesNV coopmat2_features {}; ++ coopmat2_features.pNext = nullptr; ++ coopmat2_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_2_FEATURES_NV; ++ if (coopmat2_support) { ++ last_struct->pNext = (VkBaseOutStructure *)&coopmat2_features; ++ last_struct = (VkBaseOutStructure *)&coopmat2_features; ++ device_extensions.push_back("VK_NV_cooperative_matrix2"); ++ } ++#endif ++ ++ vkGetPhysicalDeviceFeatures2(device->physical_device, &device_features2); ++ ++ device->fp16 = device->fp16 && vk12_features.shaderFloat16; ++ ++ device->pipeline_robustness = pl_robustness_features.pipelineRobustness; ++ ++ if (device->subgroup_size_control) { ++ device->subgroup_min_size = subgroup_size_control_props.minSubgroupSize; ++ device->subgroup_max_size = subgroup_size_control_props.maxSubgroupSize; ++ } ++ ++ device->subgroup_size_control = device->subgroup_size_control && ++ (subgroup_size_control_props.requiredSubgroupSizeStages & vk::ShaderStageFlagBits::eCompute) && ++ subgroup_size_control_features.subgroupSizeControl; ++ ++ if (device->subgroup_size_control) { ++ device->subgroup_require_full_support = subgroup_size_control_features.computeFullSubgroups; ++ device_extensions.push_back("VK_EXT_subgroup_size_control"); ++ } ++ ++ device->coopmat_support = device->coopmat_support && coopmat_features.cooperativeMatrix; ++ ++ if (coopmat2_support) { ++#if defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) ++ if (coopmat2_features.cooperativeMatrixWorkgroupScope && ++ coopmat2_features.cooperativeMatrixFlexibleDimensions && ++ coopmat2_features.cooperativeMatrixReductions && ++ coopmat2_features.cooperativeMatrixConversions && ++ coopmat2_features.cooperativeMatrixPerElementOperations && ++ coopmat2_features.cooperativeMatrixTensorAddressing && ++ coopmat2_features.cooperativeMatrixBlockLoads && ++ vk12_features.bufferDeviceAddress) { ++ ++ std::vector flexible_dimensions; ++ uint32_t count = 0; ++ ++ PFN_vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV ++ _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV = ++ (PFN_vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV) ++ vk_instance.instance.getProcAddr("vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV"); ++ ++ _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV(device->physical_device, &count, nullptr); ++ ++ VkCooperativeMatrixFlexibleDimensionsPropertiesNV empty_prop {}; ++ empty_prop.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_FLEXIBLE_DIMENSIONS_PROPERTIES_NV; ++ flexible_dimensions.resize(count, empty_prop); ++ ++ _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV(device->physical_device, &count, flexible_dimensions.data()); ++ ++ bool found_fp16_128 = false, ++ found_fp16_256 = false, ++ found_fp32_128 = false, ++ found_fp32_256 = false; ++ // need to support fp16*fp16 with fp16/fp32 accumulator, for workgroupsize 128 ++ // with 32x16x16 and 256 with 32x32x16. ++ for (auto &prop : flexible_dimensions) { ++ if (prop.saturatingAccumulation == VK_FALSE && ++ prop.scope == VK_SCOPE_WORKGROUP_KHR && ++ prop.AType == VK_COMPONENT_TYPE_FLOAT16_KHR && ++ prop.BType == VK_COMPONENT_TYPE_FLOAT16_KHR) { ++ ++ if (prop.workgroupInvocations == 128 && ++ prop.MGranularity <= 32 && ++ prop.NGranularity <= 16 && ++ prop.KGranularity <= 16) { ++ if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR && ++ prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) { ++ found_fp16_128 = true; ++ } ++ if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && ++ prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) { ++ found_fp32_128 = true; ++ } ++ } ++ if (prop.workgroupInvocations == 256 && ++ prop.MGranularity <= 32 && ++ prop.NGranularity <= 32 && ++ prop.KGranularity <= 16) { ++ if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR && ++ prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) { ++ found_fp16_256 = true; ++ } ++ if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && ++ prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) { ++ found_fp32_256 = true; ++ } ++ } ++ } ++ } ++ if (found_fp16_128 && found_fp16_256 && ++ found_fp32_128 && found_fp32_256 && ++ coopmat2_props.cooperativeMatrixFlexibleDimensionsMaxDimension >= 512) { ++ device->coopmat2 = true; ++ } ++ } ++#endif ++ } ++ ++ if (!vk11_features.storageBuffer16BitAccess) { ++ std::cerr << "ggml_vulkan: device " << GGML_VK_NAME << idx << " does not support 16-bit storage." << std::endl; ++ throw std::runtime_error("Unsupported device"); ++ } ++ ++ device_extensions.push_back("VK_KHR_16bit_storage"); ++ ++#ifdef GGML_VULKAN_VALIDATE ++ device_extensions.push_back("VK_KHR_shader_non_semantic_info"); ++#endif ++ ++ if (device->fp16) { ++ device_extensions.push_back("VK_KHR_shader_float16_int8"); ++ } ++ ++ if (device->coopmat_support) { ++ // Query supported shapes ++ std::vector cm_props; ++ ++ PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR = ++ (PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR)vkGetInstanceProcAddr(vk_instance.instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR"); ++ ++ uint32_t cm_props_num; ++ ++ pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR(device->physical_device, &cm_props_num, nullptr); ++ ++ cm_props.resize(cm_props_num); ++ ++ for (auto& prop : cm_props) { ++ prop.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR; ++ } ++ ++ pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR(device->physical_device, &cm_props_num, cm_props.data()); ++ ++ VK_LOG_DEBUG("ggml_vulkan: Cooperative Matrix Shapes: " << cm_props.size()); ++ ++ for (auto& prop : cm_props) { ++ VK_LOG_DEBUG("ggml_vulkan: M: " << prop.MSize << " N: " << prop.NSize << " K: " << prop.KSize << " A: " << vk::to_string((vk::ComponentTypeKHR)prop.AType) << " B: " << vk::to_string((vk::ComponentTypeKHR)prop.BType) << " C: " << vk::to_string((vk::ComponentTypeKHR)prop.CType) << " Result: " << vk::to_string((vk::ComponentTypeKHR)prop.ResultType) << " saturatingAccumulation: " << prop.saturatingAccumulation << " scope: " << vk::to_string((vk::ScopeKHR)prop.scope)); ++ ++ if ((vk::ComponentTypeKHR)prop.AType == vk::ComponentTypeKHR::eFloat16 && ++ (vk::ComponentTypeKHR)prop.BType == vk::ComponentTypeKHR::eFloat16 && ++ (vk::ScopeKHR)prop.scope == vk::ScopeKHR::eSubgroup ++ ) { ++ if ((vk::ComponentTypeKHR)prop.CType == vk::ComponentTypeKHR::eFloat32 && ++ (vk::ComponentTypeKHR)prop.ResultType == vk::ComponentTypeKHR::eFloat32) { ++ // coopmat sizes not set yet ++ if (device->coopmat_m == 0) { ++ device->coopmat_acc_f32_support = true; ++ device->coopmat_m = prop.MSize; ++ device->coopmat_n = prop.NSize; ++ device->coopmat_k = prop.KSize; ++ } else if (device->coopmat_m == prop.MSize && device->coopmat_n == prop.NSize && device->coopmat_k == prop.KSize) { ++ // Only enable if shape is identical ++ device->coopmat_acc_f32_support = true; ++ } ++ } else if ((vk::ComponentTypeKHR)prop.CType == vk::ComponentTypeKHR::eFloat16 && ++ (vk::ComponentTypeKHR)prop.ResultType == vk::ComponentTypeKHR::eFloat16) { ++ // coopmat sizes not set yet ++ if (device->coopmat_m == 0) { ++ device->coopmat_acc_f16_support = true; ++ device->coopmat_m = prop.MSize; ++ device->coopmat_n = prop.NSize; ++ device->coopmat_k = prop.KSize; ++ } else if (device->coopmat_m == prop.MSize && device->coopmat_n == prop.NSize && device->coopmat_k == prop.KSize) { ++ // Only enable if shape is identical ++ device->coopmat_acc_f16_support = true; ++ } ++ } ++ } ++ } ++ ++ if (device->coopmat_m == 0 || !device->coopmat_acc_f32_support) { ++ // No suitable matmul mode found ++ GGML_LOG_DEBUG("ggml_vulkan: WARNING: No suitable matrix core mode found. Disabling matrix cores.\n"); ++ device->coopmat_support = false; ++ } ++ } ++ ++ if (device->coopmat_support) { ++ device_extensions.push_back("VK_KHR_cooperative_matrix"); ++ } ++ ++ device->name = GGML_VK_NAME + std::to_string(idx); ++ ++ device_create_info = { ++ vk::DeviceCreateFlags(), ++ device_queue_create_infos, ++ {}, ++ device_extensions ++ }; ++ device_create_info.setPNext(&device_features2); ++ device->device = device->physical_device.createDevice(device_create_info); ++ ++ // Queues ++ ggml_vk_create_queue(device, device->compute_queue, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false); ++ ++ // Shaders ++ // Disable matmul tile sizes early if performance low or not supported ++ switch (device->vendor_id) { ++#ifndef GGML_VULKAN_RUN_TESTS ++ case VK_VENDOR_ID_AMD: ++ case VK_VENDOR_ID_INTEL: ++ device->mul_mat_l = false; ++ device->mul_mat_m = true; ++ device->mul_mat_s = true; ++ device->mul_mat_id_l = false; ++ device->mul_mat_id_m = true; ++ device->mul_mat_id_s = true; ++ break; ++ case VK_VENDOR_ID_APPLE: ++ device->mul_mat_l = false; ++ device->mul_mat_m = true; ++ device->mul_mat_s = false; ++ device->mul_mat_id_l = false; ++ device->mul_mat_id_m = true; ++ device->mul_mat_id_s = false; ++ break; ++#endif ++ default: ++ device->mul_mat_l = true; ++ device->mul_mat_m = true; ++ device->mul_mat_s = true; ++ device->mul_mat_id_l = true; ++ device->mul_mat_id_m = true; ++ device->mul_mat_id_s = true; ++ break; ++ } ++ ++ ggml_vk_load_shaders(device); ++ ++ if (!device->single_queue) { ++ const uint32_t transfer_queue_index = compute_queue_family_index == transfer_queue_family_index ? 1 : 0; ++ ggml_vk_create_queue(device, device->transfer_queue, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true); ++ } else { ++ // TODO: Use pointer or reference to avoid copy ++ device->transfer_queue = device->compute_queue; ++ } ++ ++ device->buffer_type = { ++ /* .iface = */ ggml_backend_vk_buffer_type_interface, ++ /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), idx), ++ /* .context = */ new ggml_backend_vk_buffer_type_context{ device->name, device }, ++ }; ++ ++ device->fence = device->device.createFence({}); ++ ++ device->idx = idx; ++ ++ return device; ++ } ++ ++ return vk_instance.devices[idx]; ++} ++ ++static void ggml_vk_print_gpu_info(size_t idx) { ++ GGML_ASSERT(idx < vk_instance.device_indices.size()); ++ size_t dev_num = vk_instance.device_indices[idx]; ++ VK_LOG_DEBUG("ggml_vk_print_gpu_info(" << dev_num << ")"); ++ GGML_ASSERT(vk_instance_initialized); ++ ++ std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); ++ ++ if (dev_num >= devices.size()) { ++ std::cerr << "ggml_vulkan: Device with index " << dev_num << " does not exist." << std::endl; ++ throw std::runtime_error("Device not found"); ++ } ++ ++ vk::PhysicalDevice physical_device = devices[dev_num]; ++ std::vector ext_props = physical_device.enumerateDeviceExtensionProperties(); ++ ++ vk::PhysicalDeviceProperties2 props2; ++ vk::PhysicalDeviceMaintenance3Properties props3; ++ vk::PhysicalDeviceSubgroupProperties subgroup_props; ++ vk::PhysicalDeviceDriverProperties driver_props; ++ props2.pNext = &props3; ++ props3.pNext = &subgroup_props; ++ subgroup_props.pNext = &driver_props; ++ physical_device.getProperties2(&props2); ++ ++ const size_t subgroup_size = subgroup_props.subgroupSize; ++ const bool uma = props2.properties.deviceType == vk::PhysicalDeviceType::eIntegratedGpu; ++ ++ bool fp16_storage = false; ++ bool fp16_compute = false; ++ bool coopmat_support = false; ++ bool coopmat2_support = false; ++ ++ for (auto properties : ext_props) { ++ if (strcmp("VK_KHR_16bit_storage", properties.extensionName) == 0) { ++ fp16_storage = true; ++ } else if (strcmp("VK_KHR_shader_float16_int8", properties.extensionName) == 0) { ++ fp16_compute = true; ++ } else if (strcmp("VK_KHR_cooperative_matrix", properties.extensionName) == 0 && ++ !getenv("GGML_VK_DISABLE_COOPMAT")) { ++ coopmat_support = true; ++#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) ++ } else if (strcmp("VK_NV_cooperative_matrix2", properties.extensionName) == 0 && ++ !getenv("GGML_VK_DISABLE_COOPMAT2")) { ++ coopmat2_support = true; ++#endif ++ } ++ } ++ ++ if (!ggml_vk_khr_cooperative_matrix_support(props2.properties, driver_props)) { ++ coopmat_support = false; ++ } ++ ++ const char* GGML_VK_DISABLE_F16 = getenv("GGML_VK_DISABLE_F16"); ++ bool force_disable_f16 = GGML_VK_DISABLE_F16 != nullptr; ++ ++ bool fp16 = !force_disable_f16 && fp16_storage && fp16_compute; ++ ++ vk::PhysicalDeviceFeatures device_features = physical_device.getFeatures(); ++ ++ VkPhysicalDeviceFeatures2 device_features2; ++ device_features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2; ++ device_features2.pNext = nullptr; ++ device_features2.features = (VkPhysicalDeviceFeatures)device_features; ++ ++ VkPhysicalDeviceVulkan11Features vk11_features; ++ vk11_features.pNext = nullptr; ++ vk11_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES; ++ device_features2.pNext = &vk11_features; ++ ++ VkPhysicalDeviceVulkan12Features vk12_features; ++ vk12_features.pNext = nullptr; ++ vk12_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES; ++ vk11_features.pNext = &vk12_features; ++ ++ // Pointer to the last chain element ++ VkBaseOutStructure * last_struct = (VkBaseOutStructure *)&vk12_features; ++ ++ VkPhysicalDeviceCooperativeMatrixFeaturesKHR coopmat_features; ++ coopmat_features.pNext = nullptr; ++ coopmat_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR; ++ coopmat_features.cooperativeMatrix = VK_FALSE; ++ ++ if (coopmat_support) { ++ last_struct->pNext = (VkBaseOutStructure *)&coopmat_features; ++ last_struct = (VkBaseOutStructure *)&coopmat_features; ++ } ++ ++ vkGetPhysicalDeviceFeatures2(physical_device, &device_features2); ++ ++ fp16 = fp16 && vk12_features.shaderFloat16; ++ ++ coopmat_support = coopmat_support && coopmat_features.cooperativeMatrix; ++ ++ std::string matrix_cores = coopmat2_support ? "NV_coopmat2" : coopmat_support ? "KHR_coopmat" : "none"; ++ ++ std::string device_name = props2.properties.deviceName.data(); ++ GGML_LOG_DEBUG("ggml_vulkan: %zu = %s (%s) | uma: %d | fp16: %d | warp size: %zu | matrix cores: %s\n", ++ idx, device_name.c_str(), driver_props.driverName.data(), uma, fp16, subgroup_size, matrix_cores.c_str()); ++ ++ if (props2.properties.deviceType == vk::PhysicalDeviceType::eCpu) { ++ GGML_LOG_DEBUG("ggml_vulkan: Warning: Device type is CPU. This is probably not the device you want.\n"); ++ } ++} ++ ++static bool ggml_vk_instance_validation_ext_available(const std::vector& instance_extensions); ++static bool ggml_vk_instance_portability_enumeration_ext_available(const std::vector& instance_extensions); ++ ++void ggml_vk_instance_init() { ++ if (vk_instance_initialized) { ++ return; ++ } ++ VK_LOG_DEBUG("ggml_vk_instance_init()"); ++ ++ vk_instance_initialized = true; ++ ++ vk::ApplicationInfo app_info{ "ggml-vulkan", 1, nullptr, 0, VK_API_VERSION }; ++ ++ const std::vector instance_extensions = vk::enumerateInstanceExtensionProperties(); ++ const bool validation_ext = ggml_vk_instance_validation_ext_available(instance_extensions); ++#ifdef __APPLE__ ++ const bool portability_enumeration_ext = ggml_vk_instance_portability_enumeration_ext_available(instance_extensions); ++#endif ++ ++ std::vector layers; ++ ++ if (validation_ext) { ++ layers.push_back("VK_LAYER_KHRONOS_validation"); ++ } ++ std::vector extensions; ++ if (validation_ext) { ++ extensions.push_back("VK_EXT_validation_features"); ++ } ++#ifdef __APPLE__ ++ if (portability_enumeration_ext) { ++ extensions.push_back("VK_KHR_portability_enumeration"); ++ } ++#endif ++ vk::InstanceCreateInfo instance_create_info(vk::InstanceCreateFlags{}, &app_info, layers, extensions); ++#ifdef __APPLE__ ++ if (portability_enumeration_ext) { ++ instance_create_info.flags |= vk::InstanceCreateFlagBits::eEnumeratePortabilityKHR; ++ } ++#endif ++ ++ std::vector features_enable; ++ vk::ValidationFeaturesEXT validation_features; ++ ++ if (validation_ext) { ++ features_enable = { vk::ValidationFeatureEnableEXT::eBestPractices }; ++ validation_features = { ++ features_enable, ++ {}, ++ }; ++ validation_features.setPNext(nullptr); ++ instance_create_info.setPNext(&validation_features); ++ GGML_LOG_DEBUG("ggml_vulkan: Validation layers enabled\n"); ++ } ++ vk_instance.instance = vk::createInstance(instance_create_info); ++ ++ size_t num_available_devices = vk_instance.instance.enumeratePhysicalDevices().size(); ++ ++ // Emulate behavior of CUDA_VISIBLE_DEVICES for Vulkan ++ char * devices_env = getenv("GGML_VK_VISIBLE_DEVICES"); ++ if (devices_env != nullptr) { ++ std::string devices(devices_env); ++ std::replace(devices.begin(), devices.end(), ',', ' '); ++ ++ std::stringstream ss(devices); ++ size_t tmp; ++ while (ss >> tmp) { ++ if(tmp >= num_available_devices) { ++ std::cerr << "ggml_vulkan: Invalid device index " << tmp << " in GGML_VK_VISIBLE_DEVICES." << std::endl; ++ throw std::runtime_error("Invalid Vulkan device index"); ++ } ++ vk_instance.device_indices.push_back(tmp); ++ } ++ } else { ++ std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); ++ ++ // Make sure at least one device exists ++ if (devices.empty()) { ++ std::cerr << "ggml_vulkan: Error: No devices found." << std::endl; ++ GGML_ABORT("fatal error"); ++ } ++ ++ // Default to using all dedicated GPUs ++ for (size_t i = 0; i < devices.size(); i++) { ++ vk::PhysicalDeviceProperties2 new_props; ++ vk::PhysicalDeviceDriverProperties new_driver; ++ vk::PhysicalDeviceIDProperties new_id; ++ new_props.pNext = &new_driver; ++ new_driver.pNext = &new_id; ++ devices[i].getProperties2(&new_props); ++ ++ if (new_props.properties.deviceType == vk::PhysicalDeviceType::eDiscreteGpu) { ++ // Check if there are two physical devices corresponding to the same GPU ++ auto old_device = std::find_if( ++ vk_instance.device_indices.begin(), ++ vk_instance.device_indices.end(), ++ [&devices, &new_id](const size_t k){ ++ vk::PhysicalDeviceProperties2 old_props; ++ vk::PhysicalDeviceIDProperties old_id; ++ old_props.pNext = &old_id; ++ devices[k].getProperties2(&old_props); ++ return std::equal(std::begin(old_id.deviceUUID), std::end(old_id.deviceUUID), std::begin(new_id.deviceUUID)); ++ } ++ ); ++ if (old_device == vk_instance.device_indices.end()) { ++ vk_instance.device_indices.push_back(i); ++ } else { ++ // There can be two physical devices corresponding to the same GPU if there are 2 different drivers ++ // This can cause error when splitting layers aross the devices, need to keep only 1 ++ VK_LOG_DEBUG("Device " << i << " and device " << *old_device << " have the same deviceUUID"); ++ ++ vk::PhysicalDeviceProperties2 old_props; ++ vk::PhysicalDeviceDriverProperties old_driver; ++ old_props.pNext = &old_driver; ++ devices[*old_device].getProperties2(&old_props); ++ ++ std::map driver_priorities {}; ++ int old_priority = std::numeric_limits::max(); ++ int new_priority = std::numeric_limits::max(); ++ ++ // Check https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/VkDriverId.html for the list of driver id ++ // Smaller number -> higher priority ++ switch (old_props.properties.vendorID) { ++ case VK_VENDOR_ID_AMD: ++ driver_priorities[vk::DriverId::eMesaRadv] = 1; ++ driver_priorities[vk::DriverId::eAmdOpenSource] = 2; ++ driver_priorities[vk::DriverId::eAmdProprietary] = 3; ++ break; ++ case VK_VENDOR_ID_INTEL: ++ driver_priorities[vk::DriverId::eIntelOpenSourceMESA] = 1; ++ driver_priorities[vk::DriverId::eIntelProprietaryWindows] = 2; ++ break; ++ case VK_VENDOR_ID_NVIDIA: ++ driver_priorities[vk::DriverId::eNvidiaProprietary] = 1; ++#if defined(VK_API_VERSION_1_3) && VK_HEADER_VERSION >= 235 ++ driver_priorities[vk::DriverId::eMesaNvk] = 2; ++#endif ++ break; ++ } ++ ++ if (driver_priorities.count(old_driver.driverID)) { ++ old_priority = driver_priorities[old_driver.driverID]; ++ } ++ if (driver_priorities.count(new_driver.driverID)) { ++ new_priority = driver_priorities[new_driver.driverID]; ++ } ++ ++ if (new_priority < old_priority) { ++ auto r = std::remove(vk_instance.device_indices.begin(), vk_instance.device_indices.end(), *old_device); ++ vk_instance.device_indices.erase(r, vk_instance.device_indices.end()); ++ vk_instance.device_indices.push_back(i); ++ ++ VK_LOG_DEBUG("Prioritize device " << i << " driver " << new_driver.driverName << " over device " << *old_device << " driver " << old_driver.driverName); ++ } ++ else { ++ VK_LOG_DEBUG("Prioritize device " << *old_device << " driver " << old_driver.driverName << " over device " << i << " driver " << new_driver.driverName << std::endl); ++ } ++ } ++ } ++ } ++ ++ // If no dedicated GPUs found, fall back to GPU 0 ++ if (vk_instance.device_indices.empty()) { ++ vk_instance.device_indices.push_back(0); ++ } ++ } ++ GGML_LOG_DEBUG("ggml_vulkan: Found %zu Vulkan devices:\n", vk_instance.device_indices.size()); ++ ++ for (size_t i = 0; i < vk_instance.device_indices.size(); i++) { ++ ggml_vk_print_gpu_info(i); ++ } ++} ++ ++static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) { ++ VK_LOG_DEBUG("ggml_vk_init(" << ctx->name << ", " << idx << ")"); ++ ggml_vk_instance_init(); ++ GGML_ASSERT(idx < vk_instance.device_indices.size()); ++ ++ ctx->name = GGML_VK_NAME + std::to_string(idx); ++ ++ ctx->device = ggml_vk_get_device(idx); ++ ++ ctx->semaphore_idx = 0; ++ ctx->event_idx = 0; ++ ++ ctx->prealloc_size_x = 0; ++ ctx->prealloc_size_y = 0; ++ ctx->prealloc_size_split_k = 0; ++ ++ ctx->fence = ctx->device->device.createFence({}); ++ ++#ifdef GGML_VULKAN_CHECK_RESULTS ++ const char* skip_checks = getenv("GGML_VULKAN_SKIP_CHECKS"); ++ vk_skip_checks = (skip_checks == NULL ? 0 : atoi(skip_checks)); ++ const char* output_tensor = getenv("GGML_VULKAN_OUTPUT_TENSOR"); ++ vk_output_tensor = (output_tensor == NULL ? 0 : atoi(output_tensor)); ++#endif ++} ++ ++static vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type type) { ++ VK_LOG_DEBUG("ggml_vk_get_to_fp16()"); ++ switch (type) { ++ case GGML_TYPE_F32: ++ case GGML_TYPE_Q4_0: ++ case GGML_TYPE_Q4_1: ++ case GGML_TYPE_Q5_0: ++ case GGML_TYPE_Q5_1: ++ case GGML_TYPE_Q8_0: ++ case GGML_TYPE_Q2_K: ++ case GGML_TYPE_Q3_K: ++ case GGML_TYPE_Q4_K: ++ case GGML_TYPE_Q5_K: ++ case GGML_TYPE_Q6_K: ++ case GGML_TYPE_IQ4_NL: ++ break; ++ default: ++ return nullptr; ++ } ++ ++ return ctx->device->pipeline_dequant[type]; ++} ++ ++static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec) { ++ VK_LOG_DEBUG("ggml_vk_get_mul_mat_mat_pipeline(" << ggml_type_name(src0_type) << ", " << ggml_type_name(src1_type) << ")"); ++ if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_matmul_f32; ++ } ++ if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F16) { ++ return ctx->device->pipeline_matmul_f32_f16; ++ } ++ if (prec == GGML_PREC_DEFAULT && ctx->device->fp16 && !(ctx->device->coopmat_support && !ctx->device->coopmat_acc_f16_support)) { ++ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_matmul_f16_f32.f16acc; ++ } ++ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { ++ return ctx->device->pipeline_matmul_f16.f16acc; ++ } ++ } else { ++ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_matmul_f16_f32.f32acc; ++ } ++ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { ++ return ctx->device->pipeline_matmul_f16.f32acc; ++ } ++ } ++ ++ if (src1_type != GGML_TYPE_F32 && !ctx->device->coopmat2) { ++ return nullptr; ++ } ++ ++ switch (src0_type) { ++ case GGML_TYPE_Q4_0: ++ case GGML_TYPE_Q4_1: ++ case GGML_TYPE_Q5_0: ++ case GGML_TYPE_Q5_1: ++ case GGML_TYPE_Q8_0: ++ case GGML_TYPE_Q2_K: ++ case GGML_TYPE_Q3_K: ++ case GGML_TYPE_Q4_K: ++ case GGML_TYPE_Q5_K: ++ case GGML_TYPE_Q6_K: ++ case GGML_TYPE_IQ4_NL: ++ break; ++ default: ++ return nullptr; ++ } ++ ++ if (ctx->device->coopmat2) { ++ assert(src1_type == GGML_TYPE_F16); ++ return ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f16acc; ++ } ++ return ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f32acc; ++} ++ ++static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * ctx, ggml_type a_type, ggml_type b_type, uint32_t num_cols) { ++ VK_LOG_DEBUG("ggml_vk_get_dequantize_mul_mat_vec()"); ++ GGML_ASSERT(b_type == GGML_TYPE_F32 || b_type == GGML_TYPE_F16); ++ GGML_ASSERT(num_cols >= 1 && num_cols <= mul_mat_vec_max_cols); ++ ++ switch (a_type) { ++ case GGML_TYPE_F32: ++ case GGML_TYPE_F16: ++ case GGML_TYPE_Q4_0: ++ case GGML_TYPE_Q4_1: ++ case GGML_TYPE_Q5_0: ++ case GGML_TYPE_Q5_1: ++ case GGML_TYPE_Q8_0: ++ case GGML_TYPE_Q2_K: ++ case GGML_TYPE_Q3_K: ++ case GGML_TYPE_Q4_K: ++ case GGML_TYPE_Q5_K: ++ case GGML_TYPE_Q6_K: ++ case GGML_TYPE_IQ4_NL: ++ break; ++ default: ++ return nullptr; ++ } ++ ++ return b_type == GGML_TYPE_F32 ? ctx->device->pipeline_dequant_mul_mat_vec_f32_f32[a_type][num_cols-1] : ctx->device->pipeline_dequant_mul_mat_vec_f16_f32[a_type][num_cols-1]; ++} ++ ++static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec) { ++ VK_LOG_DEBUG("ggml_vk_get_mul_mat_mat_id_pipeline()"); ++ if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_matmul_id_f32; ++ } ++ if (prec == GGML_PREC_DEFAULT && ctx->device->fp16 && !(ctx->device->coopmat_support && !ctx->device->coopmat_acc_f16_support)) { ++ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_matmul_id_f16_f32.f16acc; ++ } ++ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { ++ return ctx->device->pipeline_matmul_id_f16.f16acc; ++ } ++ } else { ++ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_matmul_id_f16_f32.f32acc; ++ } ++ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { ++ return ctx->device->pipeline_matmul_id_f16.f32acc; ++ } ++ } ++ ++ GGML_ASSERT(src1_type == GGML_TYPE_F32); ++ ++ switch (src0_type) { ++ case GGML_TYPE_Q4_0: ++ case GGML_TYPE_Q4_1: ++ case GGML_TYPE_Q5_0: ++ case GGML_TYPE_Q5_1: ++ case GGML_TYPE_Q8_0: ++ case GGML_TYPE_Q2_K: ++ case GGML_TYPE_Q3_K: ++ case GGML_TYPE_Q4_K: ++ case GGML_TYPE_Q5_K: ++ case GGML_TYPE_Q6_K: ++ case GGML_TYPE_IQ4_NL: ++ break; ++ default: ++ return nullptr; ++ } ++ ++ return ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat_id[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat_id[src0_type].f32acc; ++} ++ ++static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context * ctx, ggml_type a_type, ggml_type b_type) { ++ VK_LOG_DEBUG("ggml_vk_get_dequantize_mul_mat_vec()"); ++ GGML_ASSERT(b_type == GGML_TYPE_F32); ++ ++ switch (a_type) { ++ case GGML_TYPE_F32: ++ case GGML_TYPE_F16: ++ case GGML_TYPE_Q4_0: ++ case GGML_TYPE_Q4_1: ++ case GGML_TYPE_Q5_0: ++ case GGML_TYPE_Q5_1: ++ case GGML_TYPE_Q8_0: ++ case GGML_TYPE_Q2_K: ++ case GGML_TYPE_Q3_K: ++ case GGML_TYPE_Q4_K: ++ case GGML_TYPE_Q5_K: ++ case GGML_TYPE_Q6_K: ++ case GGML_TYPE_IQ4_NL: ++ break; ++ default: ++ return nullptr; ++ } ++ ++ return ctx->device->pipeline_dequant_mul_mat_vec_id_f32[a_type]; ++} ++ ++static vk_buffer ggml_vk_pool_malloc(ggml_backend_vk_context * ctx, size_t size) { ++ VK_LOG_DEBUG("ggml_vk_pool_malloc(" << size << ")"); ++ VK_LOG_MEMORY("ggml_vk_pool_malloc"); ++ ++ int best_i = -1; ++ size_t best_size = std::numeric_limits::max(); //smallest unused buffer that fits our needs ++ int worst_i = -1; ++ size_t worst_size = 0; //largest unused buffer seen so far ++ for (int i = 0; i < MAX_VK_BUFFERS; ++i) { ++ vk_buffer &b = ctx->buffer_pool[i]; ++ if (b != nullptr && b->size >= size && b->size < best_size) { ++ best_i = i; ++ best_size = b->size; ++ } ++ if (b != nullptr && b->size > worst_size) { ++ worst_i = i; ++ worst_size = b->size; ++ } ++ } ++ if(best_i != -1) { ++ //found the smallest buffer that fits our needs ++ vk_buffer b = ctx->buffer_pool[best_i]; ++ ctx->buffer_pool[best_i].reset(); ++ return b; ++ } ++ if(worst_i != -1) { ++ //no buffer that fits our needs, resize largest one to save memory ++ vk_buffer& b = ctx->buffer_pool[worst_i]; ++ ggml_vk_destroy_buffer(b); ++ } ++ ++ return ggml_vk_create_buffer_device(ctx->device, size); ++} ++ ++static void ggml_vk_pool_free(ggml_backend_vk_context * ctx, vk_buffer& buffer) { ++ VK_LOG_DEBUG("ggml_vk_pool_free(" << buffer->size << ")"); ++ for (int i = 0; i < MAX_VK_BUFFERS; ++i) { ++ vk_buffer& b = ctx->buffer_pool[i]; ++ if (b == nullptr) { ++ b = buffer; ++ return; ++ } ++ } ++ std::cerr << "ggml_vulkan: WARNING: vk buffer pool full, increase MAX_VK_BUFFERS" << std::endl; ++ ggml_vk_destroy_buffer(buffer); ++} ++ ++// Returns an available temporary buffer that may only be used temporarily, it will be reused ++static vk_buffer ggml_vk_create_buffer_temp(ggml_backend_vk_context * ctx, size_t size) { ++ // Try to find existing temp buffer with enough capacity ++ for (auto& buffer : ctx->gc.temp_buffers) { ++ if (buffer->size >= size) { ++ return buffer; ++ } ++ } ++ ++ VK_LOG_MEMORY("ggml_vk_create_buffer_temp(" << size << ")"); ++ ++ // Otherwise create new buffer ++ vk_buffer buf = ggml_vk_pool_malloc(ctx, size); ++ ctx->gc.temp_buffers.push_back(buf); ++ ++ return buf; ++} ++ ++static void * ggml_vk_host_malloc(vk_device& device, size_t size) { ++ VK_LOG_MEMORY("ggml_vk_host_malloc(" << size << ")"); ++ vk_buffer buf = ggml_vk_create_buffer(device, size, ++ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached, ++ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); ++ ++ if(!(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible)) { ++ fprintf(stderr, "WARNING: failed to allocate %.2f MB of pinned memory\n", ++ size/1024.0/1024.0); ++ device->device.freeMemory(buf->device_memory); ++ device->device.destroyBuffer(buf->buffer); ++ return nullptr; ++ } ++ ++ device->pinned_memory.push_back(std::make_tuple(buf->ptr, size, buf)); ++ ++ return buf->ptr; ++} ++ ++static void ggml_vk_host_free(vk_device& device, void* ptr) { ++ if (ptr == nullptr) { ++ return; ++ } ++ VK_LOG_MEMORY("ggml_vk_host_free(" << ptr << ")"); ++ vk_buffer buf; ++ size_t index; ++ for (size_t i = 0; i < device->pinned_memory.size(); i++) { ++ const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]); ++ const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]); ++ if (ptr >= addr && ptr < endr) { ++ buf = std::get<2>(device->pinned_memory[i]); ++ index = i; ++ break; ++ } ++ } ++ if (buf == nullptr) { ++ fprintf(stderr, "WARNING: failed to free pinned memory: memory not in map\n"); ++ return; ++ } ++ ++ ggml_vk_destroy_buffer(buf); ++ ++ device->pinned_memory.erase(device->pinned_memory.begin() + index); ++} ++ ++static void ggml_vk_host_get(vk_device& device, const void * ptr, vk_buffer& buf, size_t& buf_offset) { ++ buf = nullptr; ++ buf_offset = 0; ++ for (size_t i = 0; i < device->pinned_memory.size(); i++) { ++ const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]); ++ const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]); ++ if (ptr >= addr && ptr < endr) { ++ buf = std::get<2>(device->pinned_memory[i]); ++ buf_offset = ((const uint8_t *)ptr) - addr; ++ break; ++ } ++ } ++} ++ ++static vk_submission ggml_vk_begin_submission(vk_device& device, vk_queue& q, bool one_time = true) { ++ vk_submission s; ++ s.buffer = ggml_vk_create_cmd_buffer(device, q); ++ if (one_time) { ++ s.buffer.begin({ vk::CommandBufferUsageFlagBits::eOneTimeSubmit }); ++ } else { ++ s.buffer.begin({ vk::CommandBufferUsageFlags{} }); ++ } ++ ++ return s; ++} ++ ++ ++ ++static void ggml_vk_dispatch_pipeline(ggml_backend_vk_context* ctx, vk_context& subctx, vk_pipeline& pipeline, std::initializer_list const& descriptor_buffer_infos, size_t push_constant_size, const void* push_constants, std::array elements) { ++ const uint32_t wg0 = CEIL_DIV(elements[0], pipeline->wg_denoms[0]); ++ const uint32_t wg1 = CEIL_DIV(elements[1], pipeline->wg_denoms[1]); ++ const uint32_t wg2 = CEIL_DIV(elements[2], pipeline->wg_denoms[2]); ++ VK_LOG_DEBUG("ggml_vk_dispatch_pipeline(" << pipeline->name << ", {"; ++ for (auto& buffer : descriptor_buffer_infos) { ++ std::cerr << "(" << buffer.buffer << ", " << buffer.offset << ", " << buffer.range << "), "; ++ } ++ std::cerr << "}, (" << wg0 << "," << wg1 << "," << wg2 << "))"); ++ GGML_ASSERT(pipeline->descriptor_set_idx < pipeline->descriptor_sets.size()); ++ GGML_ASSERT(descriptor_buffer_infos.size() == pipeline->parameter_count); ++ ++ vk::DescriptorSet& descriptor_set = pipeline->descriptor_sets[pipeline->descriptor_set_idx++]; ++ vk::WriteDescriptorSet write_descriptor_set{ descriptor_set, 0, 0, pipeline->parameter_count, vk::DescriptorType::eStorageBuffer, nullptr, descriptor_buffer_infos.begin() }; ++ ctx->device->device.updateDescriptorSets({ write_descriptor_set }, {}); ++ ++ subctx->s->buffer.pushConstants(pipeline->layout, vk::ShaderStageFlagBits::eCompute, 0, push_constant_size, push_constants); ++ subctx->s->buffer.bindPipeline(vk::PipelineBindPoint::eCompute, pipeline->pipeline); ++ subctx->s->buffer.bindDescriptorSets(vk::PipelineBindPoint::eCompute, ++ pipeline->layout, ++ 0, ++ { descriptor_set }, ++ {}); ++ subctx->s->buffer.dispatch(wg0, wg1, wg2); ++} ++ ++static void ggml_vk_end_submission(vk_submission& s, std::vector wait_semaphores, std::vector signal_semaphores) { ++ s.buffer.end(); ++ ++ s.wait_semaphores = std::move(wait_semaphores); ++ s.signal_semaphores = std::move(signal_semaphores); ++} ++ ++static void ggml_vk_ctx_end(vk_context& ctx) { ++ VK_LOG_DEBUG("ggml_vk_ctx_end(" << ctx << ", " << ctx->seqs.size() << ")"); ++ if (ctx->s == nullptr) { ++ return; ++ } ++ ++ ctx->s->buffer.end(); ++ ctx->s = nullptr; ++} ++ ++static void ggml_vk_ctx_begin(vk_device& device, vk_context& subctx) { ++ VK_LOG_DEBUG("ggml_vk_ctx_begin(" << device->name << ")"); ++ if (subctx->s != nullptr) { ++ ggml_vk_ctx_end(subctx); ++ } ++ ++ subctx->seqs.push_back({ ggml_vk_begin_submission(device, *subctx->q) }); ++ subctx->s = subctx->seqs[subctx->seqs.size() - 1].data(); ++} ++ ++static size_t ggml_vk_align_size(size_t width, size_t align) { ++ VK_LOG_DEBUG("ggml_vk_align_size(" << width << ", " << align << ")"); ++ return CEIL_DIV(width, align) * align; ++} ++ ++static void deferred_memcpy(void * dst, const void * src, size_t size, std::vector* memcpys = nullptr) { ++ if (memcpys == nullptr) { ++ memcpy(dst, src, size); ++ } else { ++ memcpys->emplace_back(dst, src, size); ++ } ++} ++ ++static void ggml_vk_ensure_sync_staging_buffer(vk_device& device, size_t size) { ++ if (device->sync_staging == nullptr || device->sync_staging->size < size) { ++ VK_LOG_MEMORY("ggml_vk_ensure_sync_staging_buffer(" << size << ")"); ++ ggml_vk_destroy_buffer(device->sync_staging); ++ device->sync_staging = ggml_vk_create_buffer_check(device, size, ++ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached, ++ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); ++ } ++} ++ ++static void ggml_vk_buffer_write_nc_async(ggml_backend_vk_context * ctx, vk_context& subctx, vk_buffer& dst, size_t offset, const ggml_tensor * tensor, bool sync_staging = false) { ++ VK_LOG_DEBUG("ggml_vk_buffer_write_nc_async(" << tensor << ")"); ++ GGML_ASSERT(!ggml_is_contiguous(tensor)); ++ // Buffer is already mapped ++ if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { ++ std::cerr << "ggml_vulkan: buffer_write_nc_async dst buffer is host_visible. Use synchronous write." << std::endl; ++ GGML_ABORT("fatal error"); ++ } ++ // Check if src is pinned memory ++ vk_buffer buf = nullptr; ++ size_t buf_offset = 0; ++ ggml_vk_host_get(ctx->device, tensor->data, buf, buf_offset); ++ ++ const uint64_t ne0 = tensor->ne[0]; ++ const uint64_t ne1 = tensor->ne[1]; ++ const uint64_t ne2 = tensor->ne[2]; ++ const uint64_t ne3 = tensor->ne[3]; ++ const uint64_t nb0 = tensor->nb[0]; ++ const uint64_t nb1 = tensor->nb[1]; ++ const uint64_t nb2 = tensor->nb[2]; ++ const uint64_t nb3 = tensor->nb[3]; ++ const ggml_type type = tensor->type; ++ const uint64_t ts = ggml_type_size(type); ++ const uint64_t bs = ggml_blck_size(type); ++ ++ const uint64_t dstnb0 = ts; ++ const uint64_t dstnb1 = dstnb0*(ne0/bs); ++ const uint64_t dstnb2 = dstnb1*ne1; ++ const uint64_t dstnb3 = dstnb2*ne2; ++ ++ const uint64_t ne = ggml_nelements(tensor); ++ ++ if (buf != nullptr) { ++ // Memory is pinned, use as staging buffer ++ std::vector slices; ++ ++ for (uint64_t i3 = 0; i3 < ne3; i3++) { ++ for (uint64_t i2 = 0; i2 < ne2; i2++) { ++ // Find longest contiguous slice ++ if (ne1*nb1 == dstnb2) { ++ slices.push_back({ buf_offset + i3*nb3 + i2*nb2, offset + i3*dstnb3 + i2*dstnb2, dstnb2 }); ++ } else { ++ for (uint64_t i1 = 0; i1 < ne1; i1++) { ++ if (ne0*nb0/bs == dstnb1) { ++ slices.push_back({ buf_offset + i3*nb3 + i2*nb2 + i1*nb1, offset + i3*dstnb3 + i2*dstnb2 + i1*dstnb1, dstnb1 }); ++ } else { ++ const uint64_t s_off = buf_offset + i3*nb3 + i2*nb2 + i1*nb1; ++ const uint64_t d_off = offset + i3*dstnb3 + i2*dstnb2 + i1*dstnb1; ++ for (uint64_t i0 = 0; i0 < ne0; i0++) { ++ slices.push_back({ s_off + i1*nb0, d_off + i0*dstnb0, dstnb0 }); ++ } ++ } ++ } ++ } ++ } ++ } ++ ++ ggml_vk_sync_buffers(subctx); ++ subctx->s->buffer.copyBuffer(buf->buffer, dst->buffer, slices); ++ return; ++ } ++ ++ if (!sync_staging) { ++ GGML_ABORT("Asynchronous write to non-pinned memory not supported"); ++ } ++ ++ // Staging buffer required ++ vk_buffer& staging = ctx->device->sync_staging; ++ const uint64_t copy_size = ts*ne/bs; ++ ggml_vk_ensure_sync_staging_buffer(ctx->device, copy_size); ++ VkBufferCopy buf_copy{ 0, offset, copy_size }; ++ ++ ggml_vk_sync_buffers(subctx); ++ vkCmdCopyBuffer(subctx->s->buffer, (VkBuffer)staging->buffer, (VkBuffer)dst->buffer, 1, &buf_copy); ++ ++ for (uint64_t i3 = 0; i3 < ne3; i3++) { ++ for (uint64_t i2 = 0; i2 < ne2; i2++) { ++ // Find longest contiguous slice ++ if (ne1*nb1 == dstnb2) { ++ deferred_memcpy((uint8_t *)staging->ptr + i3*dstnb3 + i2*dstnb2, (const uint8_t *) tensor->data + buf_offset + i3*nb3 + i2*nb2, dstnb2, &subctx->in_memcpys); ++ } else { ++ for (uint64_t i1 = 0; i1 < ne1; i1++) { ++ if (ne0*nb0/bs == dstnb1) { ++ deferred_memcpy((uint8_t *)staging->ptr + i3*dstnb3 + i2*dstnb2 + i1*dstnb1, (const uint8_t *) tensor->data + buf_offset + i3*nb3 + i2*nb2 + i1*nb1, dstnb1, &subctx->in_memcpys); ++ } else { ++ const uint64_t s_off = buf_offset + i3*nb3 + i2*nb2 + i1*nb1; ++ const uint64_t d_off = i3*dstnb3 + i2*dstnb2 + i1*dstnb1; ++ for (uint64_t i0 = 0; i0 < ne0; i0++) { ++ deferred_memcpy((uint8_t *)staging->ptr + d_off + i0*dstnb0, (const uint8_t *) tensor->data + s_off + i0*nb0, dstnb0, &subctx->in_memcpys); ++ } ++ } ++ } ++ } ++ } ++ } ++} ++ ++static void ggml_vk_buffer_write_2d_async(vk_context subctx, vk_buffer& dst, size_t offset, const void * src, size_t spitch, size_t width, size_t height, bool sync_staging = false) { ++ VK_LOG_DEBUG("ggml_vk_buffer_write_2d_async(" << width << ", " << height << ")"); ++ // Buffer is already mapped ++ if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { ++ std::cerr << "ggml_vulkan: buffer_write_async dst buffer is host_visible. Use synchronous write." << std::endl; ++ GGML_ABORT("fatal error"); ++ } ++ // Check if src is pinned memory ++ vk_buffer buf = nullptr; ++ size_t buf_offset = 0; ++ ggml_vk_host_get(dst->device, src, buf, buf_offset); ++ ++ if (buf != nullptr) { ++ // Memory is pinned, use as staging buffer ++ std::vector slices(1); ++ if (width == spitch) { ++ // Only do single write if stride is equal ++ slices[0].srcOffset = buf_offset; ++ slices[0].dstOffset = offset; ++ slices[0].size = width * height; ++ } else { ++ slices.resize(height); ++ for (size_t i = 0; i < height; i++) { ++ slices[i].srcOffset = buf_offset + i * spitch; ++ slices[i].dstOffset = offset + i * width; ++ slices[i].size = width; ++ } ++ } ++ ++ ggml_vk_sync_buffers(subctx); ++ subctx->s->buffer.copyBuffer(buf->buffer, dst->buffer, slices); ++ return; ++ } ++ VK_LOG_DEBUG("STAGING"); ++ ++ if (!sync_staging) { ++ GGML_ABORT("Asynchronous write to non-pinned memory not supported"); ++ } ++ ++ // Staging buffer required ++ const size_t copy_size = width*height; ++ ggml_vk_ensure_sync_staging_buffer(dst->device, copy_size); ++ ++ vk_buffer& staging_buffer = dst->device->sync_staging; ++ ++ VkBufferCopy buf_copy = { ++ 0, ++ offset, ++ copy_size}; ++ ++ ggml_vk_sync_buffers(subctx); ++ vkCmdCopyBuffer(subctx->s->buffer, (VkBuffer)staging_buffer->buffer, (VkBuffer)dst->buffer, 1, &buf_copy); ++ ++ if (width == spitch) { ++ deferred_memcpy((uint8_t *)staging_buffer->ptr, src, width * height, &subctx->in_memcpys); ++ } else { ++ for (size_t i = 0; i < height; i++) { ++ deferred_memcpy((uint8_t *)staging_buffer->ptr + i * width, (const uint8_t *) src + i * spitch, width, &subctx->in_memcpys); ++ } ++ } ++} ++ ++static void ggml_vk_buffer_write_async(vk_context subctx, vk_buffer& dst, size_t offset, const void * src, size_t size, bool sync_staging = false) { ++ VK_LOG_DEBUG("ggml_vk_buffer_write_async(" << size << ")"); ++ return ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, size, size, 1, sync_staging); ++} ++ ++static void ggml_vk_buffer_write_2d(vk_buffer& dst, size_t offset, const void * src, size_t spitch, size_t width, size_t height) { ++ VK_LOG_DEBUG("ggml_vk_buffer_write_2d(" << width << ", " << height << ")"); ++ // Buffer is already mapped ++ if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { ++ GGML_ASSERT(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent); ++ ++ for (size_t i = 0; i < height; i++) { ++ memcpy((uint8_t *)dst->ptr + offset + i * width, (const uint8_t *) src + i * spitch, width); ++ } ++ } else { ++ vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue); ++ ggml_vk_ctx_begin(dst->device, subctx); ++ ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, spitch, width, height, true); ++ ggml_vk_ctx_end(subctx); ++ ++ for (auto& cpy : subctx->in_memcpys) { ++ memcpy(cpy.dst, cpy.src, cpy.n); ++ } ++ ++ ggml_vk_submit(subctx, dst->device->fence); ++ VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_buffer_write_2d waitForFences"); ++ dst->device->device.resetFences({ dst->device->fence }); ++ } ++} ++ ++static void ggml_vk_buffer_write(vk_buffer& dst, size_t offset, const void * src, size_t size) { ++ VK_LOG_DEBUG("ggml_vk_buffer_write(" << size << ")"); ++ ggml_vk_buffer_write_2d(dst, offset, src, 0, size, 1); ++} ++ ++static void ggml_vk_buffer_read_2d_async(vk_context subctx, vk_buffer& src, size_t offset, void * dst, size_t spitch, size_t dpitch, size_t width, size_t height, bool sync_staging = false) { ++ VK_LOG_DEBUG("ggml_vk_buffer_read_2d_async(offset=" << offset << ", width=" << width << ", height=" << height << ")"); ++ GGML_ASSERT(width > 0); ++ GGML_ASSERT(height > 0); ++ GGML_ASSERT(src != nullptr); ++ ++ // TODO: staging_offset is not used ++ ++ // Check if dst is pinned memory ++ vk_buffer buf = nullptr; ++ size_t buf_offset = 0; ++ ggml_vk_host_get(src->device, dst, buf, buf_offset); ++ ++ std::vector slices(1); ++ if (width == spitch && width == dpitch) { ++ // Only do single write if stride is equal ++ slices[0].srcOffset = offset; ++ slices[0].dstOffset = buf_offset; ++ slices[0].size = width * height; ++ } else { ++ slices.resize(height); ++ for (size_t i = 0; i < height; i++) { ++ slices[i].srcOffset = offset + i * spitch; ++ slices[i].dstOffset = buf_offset + i * dpitch; ++ slices[i].size = width; ++ } ++ } ++ ++ if (buf != nullptr) { ++ // Memory is pinned, use as staging buffer ++ ggml_vk_sync_buffers(subctx); ++ subctx->s->buffer.copyBuffer(src->buffer, buf->buffer, slices); ++ ++ return; ++ } ++ VK_LOG_DEBUG("STAGING"); ++ ++ if (!sync_staging) { ++ GGML_ABORT("Asynchronous read from non-pinned memory not supported"); ++ } ++ ++ // Fall back to staging buffer ++ const size_t copy_size = dpitch * height; ++ ggml_vk_ensure_sync_staging_buffer(src->device, copy_size); ++ ++ vk_buffer& staging_buffer = src->device->sync_staging; ++ ++ ggml_vk_sync_buffers(subctx); ++ subctx->s->buffer.copyBuffer(src->buffer, staging_buffer->buffer, slices); ++ ++ deferred_memcpy(dst, staging_buffer->ptr, copy_size, &subctx->out_memcpys); ++} ++ ++static void ggml_vk_buffer_read_async(vk_context subctx, vk_buffer& src, size_t offset, void * dst, size_t size, bool sync_staging = false) { ++ return ggml_vk_buffer_read_2d_async(subctx, src, offset, dst, size, size, size, 1, sync_staging); ++} ++ ++static void ggml_vk_buffer_read(vk_buffer& src, size_t offset, void * dst, size_t size) { ++ VK_LOG_DEBUG("ggml_vk_buffer_read(" << src->buffer << ", " << offset << ", " << size << ")"); ++ ++ // If the device is not an UMA device the memory is host-accessible through rebar. While writing ++ // through PCIe is sufficient fast reading back data from PCIe is slower than going through ++ // the HW device to host copy path. ++ if(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible && src->device->uma) { ++ GGML_ASSERT(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent); ++ ++ memcpy(dst, (uint8_t *) src->ptr + offset, size); ++ } else { ++ vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue); ++ ggml_vk_ctx_begin(src->device, subctx); ++ ggml_vk_buffer_read_async(subctx, src, offset, dst, size, true); ++ ggml_vk_ctx_end(subctx); ++ ++ ggml_vk_submit(subctx, src->device->fence); ++ VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_read waitForFences"); ++ src->device->device.resetFences({ src->device->fence }); ++ ++ for (auto& cpy : subctx->out_memcpys) { ++ memcpy(cpy.dst, cpy.src, cpy.n); ++ } ++ } ++} ++ ++static void ggml_vk_buffer_copy_async(vk_context& ctx, vk_buffer& dst, size_t dst_offset, vk_buffer& src, size_t src_offset, size_t size) { ++ VK_LOG_DEBUG("ggml_vk_buffer_copy_async(" << size << ")"); ++ // Make sure both buffers are on same device ++ GGML_ASSERT(src->device == dst->device); ++ ++ VkBufferCopy bc{ src_offset, dst_offset, size }; ++ ++ vkCmdCopyBuffer(ctx->s->buffer, (VkBuffer)src->buffer, (VkBuffer)dst->buffer, 1, &bc); ++} ++ ++static void ggml_vk_buffer_copy(vk_buffer& dst, size_t dst_offset, vk_buffer& src, size_t src_offset, size_t size) { ++ if (src->device == dst->device) { ++ VK_LOG_DEBUG("ggml_vk_buffer_copy(SINGLE_DEVICE, " << size << ")"); ++ // Copy within the device ++ vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue); ++ ggml_vk_ctx_begin(src->device, subctx); ++ ggml_vk_buffer_copy_async(subctx, dst, dst_offset, src, src_offset, size); ++ ggml_vk_ctx_end(subctx); ++ ggml_vk_submit(subctx, src->device->fence); ++ VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_copy waitForFences"); ++ src->device->device.resetFences({ src->device->fence }); ++ } else { ++ VK_LOG_DEBUG("ggml_vk_buffer_copy(MULTI_DEVICE, " << size << ")"); ++ // Copy device to device ++ ggml_vk_ensure_sync_staging_buffer(src->device, size); ++ ggml_vk_ensure_sync_staging_buffer(dst->device, size); ++ ++ // Copy to src staging buffer ++ ggml_vk_buffer_copy(src->device->sync_staging, 0, src, src_offset, size); ++ // memcpy to dst staging buffer ++ memcpy(dst->device->sync_staging->ptr, src->device->sync_staging->ptr, size); ++ // Copy to dst buffer ++ ggml_vk_buffer_copy(dst, dst_offset, dst->device->sync_staging, 0, size); ++ } ++} ++ ++static void ggml_vk_buffer_memset(vk_buffer& dst, size_t offset, uint32_t c, size_t size) { ++ VK_LOG_DEBUG("ggml_vk_buffer_memset(" << offset << ", " << c << ", " << size << ")"); ++ ++ vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue); ++ ggml_vk_ctx_begin(dst->device, subctx); ++ subctx->s->buffer.fillBuffer(dst->buffer, offset, size, c); ++ ggml_vk_ctx_end(subctx); ++ ++ ggml_vk_submit(subctx, dst->device->fence); ++ VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_memset waitForFences"); ++ dst->device->device.resetFences({ dst->device->fence }); ++} ++ ++static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, int m, int n, int k, const vk_pipeline& pipeline) { ++ VK_LOG_DEBUG("ggml_vk_guess_split_k(" << m << ", " << n << ", " << k << ")"); ++ ++ uint32_t split_k = 1; ++ if (ctx->device->shader_core_count != 0 && m >= (int)pipeline->wg_denoms[0] && n >= (int)pipeline->wg_denoms[1]) { ++ // If k is 'large' and the SMs will fill less than halfway, use split_k. ++ uint32_t m_tiles = CEIL_DIV(m, pipeline->wg_denoms[0]); ++ uint32_t n_tiles = CEIL_DIV(n, pipeline->wg_denoms[1]); ++ if (k >= 2048 && m_tiles * n_tiles < ctx->device->shader_core_count / 2) { ++ split_k = ctx->device->shader_core_count / (m_tiles * n_tiles); ++ // Clamp to 2 or 4 ++ split_k = std::min(split_k, 4u); ++ if (split_k == 3) { ++ split_k = 2; ++ } ++ } ++ } ++ ++ return split_k; ++} ++ ++static vk_pipeline ggml_vk_guess_matmul_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, bool aligned) { ++ VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline(" << m << ", " << n << ", " << aligned << ")"); ++ ++ if (ctx->device->coopmat2) { ++ if ((ctx->device->mul_mat_l && (m % mmp->l->wg_denoms[0]) == 0 && (n % mmp->l->wg_denoms[1]) == 0) || (!ctx->device->mul_mat_m && !ctx->device->mul_mat_s)) { ++ return aligned ? mmp->a_l : mmp->l; ++ } ++ if ((ctx->device->mul_mat_m && (m % mmp->m->wg_denoms[0]) == 0 && (n % mmp->m->wg_denoms[1]) == 0) || !ctx->device->mul_mat_s) { ++ return aligned ? mmp->a_m : mmp->m; ++ } ++ return aligned ? mmp->a_s : mmp->s; ++ } ++ ++ if ((ctx->device->mul_mat_s && (m <= 32 || n <= 32)) || (!ctx->device->mul_mat_m && !ctx->device->mul_mat_l)) { ++ return aligned ? mmp->a_s : mmp->s; ++ } ++ if ((ctx->device->mul_mat_m && (m <= 64 || n <= 64)) || !ctx->device->mul_mat_l) { ++ return aligned ? mmp->a_m : mmp->m; ++ } ++ return aligned ? mmp->a_l : mmp->l; ++} ++ ++static uint32_t ggml_vk_guess_matmul_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n) { ++ VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ")"); ++ return ggml_vk_guess_matmul_pipeline(ctx, mmp, m, n, true)->align; ++} ++ ++static void ggml_vk_matmul( ++ ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, ++ vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& split_k_buffer, ++ uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, ++ uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, ++ uint32_t split_k, uint32_t batch, uint32_t ne02, uint32_t ne12, uint32_t broadcast2, uint32_t broadcast3) { ++ VK_LOG_DEBUG("ggml_vk_matmul(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), split_k: (" << (split_k_buffer.buffer != nullptr ? split_k_buffer.buffer->buffer : VK_NULL_HANDLE) << ", " << split_k_buffer.offset << ", " << split_k_buffer.size << "), m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", split_k: " << split_k << ", batch: " << batch << ", ne02: " << ne02 << ", ne12: " << ne12 << ", broadcast2: " << broadcast2 << ", broadcast3: " << broadcast3 << ")"); ++ ggml_vk_sync_buffers(subctx); ++ if (split_k == 1) { ++ const vk_mat_mat_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, k, ne02, ne12, broadcast2, broadcast3 }; ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d }, sizeof(vk_mat_mat_push_constants), &pc, { m, n, batch }); ++ return; ++ } ++ ++ GGML_ASSERT(batch_stride_d == m * n); ++ ++ const vk_mat_mat_push_constants pc1 = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, CEIL_DIV(k, split_k), ne02, ne12, broadcast2, broadcast3 }; ++ // Make sure enough workgroups get assigned for split k to work ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, split_k_buffer }, sizeof(vk_mat_mat_push_constants), &pc1, { (CEIL_DIV(m, pipeline->wg_denoms[0]) * pipeline->wg_denoms[0]) * split_k, n, batch }); ++ ggml_vk_sync_buffers(subctx); ++ const std::array pc2 = { (uint32_t)(m * n * batch), split_k }; ++ ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_matmul_split_k_reduce, { split_k_buffer, d }, pc2.size() * sizeof(uint32_t), pc2.data(), { m * n * batch, 1, 1 }); ++} ++ ++static vk_pipeline ggml_vk_guess_matmul_id_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, bool aligned) { ++ VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline(" << m << ", " << n << ", " << aligned << ")"); ++ ++ if (ctx->device->coopmat2) { ++ if ((ctx->device->mul_mat_id_l && (m % mmp->l->wg_denoms[0]) == 0 && (n % mmp->l->wg_denoms[1]) == 0) || (!ctx->device->mul_mat_id_m && !ctx->device->mul_mat_id_s)) { ++ return aligned ? mmp->a_l : mmp->l; ++ } ++ if ((ctx->device->mul_mat_id_m && (m % mmp->m->wg_denoms[0]) == 0 && (n % mmp->m->wg_denoms[1]) == 0) || !ctx->device->mul_mat_id_s) { ++ return aligned ? mmp->a_m : mmp->m; ++ } ++ return aligned ? mmp->a_s : mmp->s; ++ } ++ ++ if ((ctx->device->mul_mat_id_s && (m <= 32 || n <= 32)) || (!ctx->device->mul_mat_id_m && !ctx->device->mul_mat_id_l)) { ++ return aligned ? mmp->a_s : mmp->s; ++ } ++ if ((ctx->device->mul_mat_id_m && (m <= 64 || n <= 64)) || !ctx->device->mul_mat_id_l) { ++ return aligned ? mmp->a_m : mmp->m; ++ } ++ return aligned ? mmp->a_l : mmp->l; ++} ++ ++static uint32_t ggml_vk_guess_matmul_id_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n) { ++ VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ")"); ++ return ggml_vk_guess_matmul_id_pipeline(ctx, mmp, m, n, true)->align; ++} ++ ++static void ggml_vk_matmul_id( ++ ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, ++ vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& ids, ++ uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, ++ uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, ++ uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11) { ++ VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), " << ++ "m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " << ++ "batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " << ++ "n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")"); ++ ggml_vk_sync_buffers(subctx); ++ const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, ++ nei0, nei1, nbi1, ne11 }; ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids }, sizeof(vk_mat_mat_id_push_constants), &pc, { m, nei1, n_as }); ++} ++ ++static bool ggml_vk_dim01_contiguous(const ggml_tensor * tensor) { ++ return ++ tensor->nb[0] == ggml_type_size(tensor->type) && ++ tensor->nb[1] == (tensor->nb[0]*tensor->ne[0])/ggml_blck_size(tensor->type) && ++ tensor->nb[3] == tensor->nb[2]*tensor->ne[2]; ++} ++ ++static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src, const ggml_tensor * dst, ggml_type to) { ++ ++ // Choose "contiguous copy" shader if src/dst are contiguous ++ bool contig = ggml_is_contiguous(src) && (!dst || ggml_is_contiguous(dst)); ++ ++ if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F32) { ++ if (contig) { ++ return ctx->device->pipeline_contig_cpy_f32_f32; ++ } else { ++ return ctx->device->pipeline_cpy_f32_f32; ++ } ++ } ++ if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F16) { ++ if (contig) { ++ return ctx->device->pipeline_contig_cpy_f32_f16; ++ } else { ++ return ctx->device->pipeline_cpy_f32_f16; ++ } ++ } ++ if (src->type == GGML_TYPE_F16 && to == GGML_TYPE_F16) { ++ if (contig) { ++ return ctx->device->pipeline_contig_cpy_f16_f16; ++ } else { ++ return ctx->device->pipeline_cpy_f16_f16; ++ } ++ } ++ ++ std::cerr << "Missing CPY op for types: " << ggml_type_name(src->type) << " " << ggml_type_name(to) << std::endl; ++ GGML_ABORT("fatal error"); ++} ++ ++static void ggml_vk_cpy_to_contiguous(ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline pipeline, const ggml_tensor * tensor, vk_subbuffer&& in, vk_subbuffer&& out) { ++ VK_LOG_DEBUG("ggml_vk_cpy_to_contiguous((" << tensor << ", type=" << tensor->type << ", ne0=" << tensor->ne[0] << ", ne1=" << tensor->ne[1] << ", ne2=" << tensor->ne[2] << ", ne3=" << tensor->ne[3] << ", nb0=" << tensor->nb[0] << ", nb1=" << tensor->nb[1] << ", nb2=" << tensor->nb[2] << ", nb3=" << tensor->nb[3] << "), "; ++ std::cerr << "buffer in size=" << in.buffer->size << ", buffer out size=" << out.buffer->size << ")"); ++ const int tensor_type_size = ggml_type_size(tensor->type); ++ ++ const uint32_t ne = ggml_nelements(tensor); ++ std::array elements; ++ ++ if (ne > 262144) { ++ elements = { 512, 512, CEIL_DIV(ne, 262144) }; ++ } else if (ne > 512) { ++ elements = { 512, CEIL_DIV(ne, 512), 1 }; ++ } else { ++ elements = { ne, 1, 1 }; ++ } ++ ++ vk_op_unary_push_constants pc = { ++ (uint32_t)ne, ++ (uint32_t)tensor->ne[0], (uint32_t)tensor->ne[1], (uint32_t)tensor->ne[2], (uint32_t)tensor->ne[3], (uint32_t)tensor->nb[0] / tensor_type_size, (uint32_t)tensor->nb[1] / tensor_type_size, (uint32_t)tensor->nb[2] / tensor_type_size, (uint32_t)tensor->nb[3] / tensor_type_size, ++ (uint32_t)tensor->ne[0], (uint32_t)tensor->ne[1], (uint32_t)tensor->ne[2], (uint32_t)tensor->ne[3], 1 , (uint32_t)tensor->ne[0] , (uint32_t)(tensor->ne[0] * tensor->ne[1]) , (uint32_t)(tensor->ne[0] * tensor->ne[1] * tensor->ne[2]), ++ 0, ++ 0.0f, 0.0f, ++ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ++ }; ++ init_pushconst_fastdiv(pc); ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { in, out }, sizeof(vk_op_unary_push_constants), &pc, elements); ++} ++ ++static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ VK_LOG_DEBUG("ggml_vk_mul_mat_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; ++ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; ++ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; ++ std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); ++ GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT ++ GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT ++ ++ const uint64_t ne00 = src0->ne[0]; ++ const uint64_t ne01 = src0->ne[1]; ++ const uint64_t ne02 = src0->ne[2]; ++ const uint64_t ne03 = src0->ne[3]; ++ ++ const uint64_t ne10 = src1->ne[0]; ++ const uint64_t ne11 = src1->ne[1]; ++ const uint64_t ne12 = src1->ne[2]; ++ const uint64_t ne13 = src1->ne[3]; ++ ++ const uint64_t ne20 = dst->ne[0]; ++ const uint64_t ne21 = dst->ne[1]; ++ ++ const uint64_t r2 = ne12 / ne02; ++ const uint64_t r3 = ne13 / ne03; ++ ++ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ++ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; ++ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; ++ ++ vk_buffer d_Qx = nullptr; ++ size_t qx_buf_offset = 0; ++ vk_buffer d_Qy = nullptr; ++ size_t qy_buf_offset = 0; ++ ++ bool src0_uma = false; ++ bool src1_uma = false; ++ ++ if (ctx->device->uma) { ++ ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); ++ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); ++ src0_uma = d_Qx != nullptr; ++ src1_uma = d_Qy != nullptr; ++ } ++ ++ const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); ++ // Reformat and convert to fp16 if src1 is non-contiguous, or for coopmat2 for better perf ++ const bool y_non_contig = (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || ++ !ggml_vk_dim01_contiguous(src1); ++ ++ const bool y_f32_kernel = src1->type == GGML_TYPE_F32 && !y_non_contig; ++ ++ vk_matmul_pipeline mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, src0->type, y_non_contig ? GGML_TYPE_F16 : src1->type, (ggml_prec)dst->op_params[0]); ++ ++ const bool qx_needs_dequant = mmp == nullptr || x_non_contig; ++ const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !y_f32_kernel) || y_non_contig; ++ ++ if (qx_needs_dequant) { ++ // Fall back to dequant + f16 mulmat ++ mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, GGML_TYPE_F16, y_f32_kernel ? GGML_TYPE_F32 : GGML_TYPE_F16, (ggml_prec)dst->op_params[0]); ++ } ++ ++ // Not implemented ++ GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT ++ ++ const int x_ne = ne01 * ne00; ++ const int y_ne = ne11 * ne10; ++ const int d_ne = ne11 * ne01; ++ ++ const uint32_t kpad = ggml_vk_align_size(ne10, ggml_vk_guess_matmul_pipeline_align(ctx, mmp, ne01, ne11)); ++ const bool aligned = ne10 == kpad && ne01 > 8 && ne11 > 8; ++ ++ vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline(ctx, mmp, ne01, ne11, aligned); ++ ++ const uint32_t split_k = ggml_vk_guess_split_k(ctx, ne01, ne11, ne10, pipeline); ++ ++ const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type); ++ const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); ++ const uint64_t x_sz = !qx_needs_dequant ? qx_sz : sizeof(ggml_fp16_t) * x_ne; ++ const uint64_t y_sz = y_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; ++ const uint64_t d_sz = sizeof(float) * d_ne; ++ ++ vk_pipeline to_fp16_vk_0 = nullptr; ++ vk_pipeline to_fp16_vk_1 = nullptr; ++ ++ if (x_non_contig) { ++ to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, GGML_TYPE_F16); ++ } else { ++ to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, src0->type); ++ } ++ if (y_non_contig) { ++ to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, GGML_TYPE_F16); ++ } else { ++ to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); ++ } ++ GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT ++ GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT ++ ++ if (dryrun) { ++ const uint64_t x_sz_upd = x_sz * ne02 * ne03; ++ const uint64_t y_sz_upd = y_sz * ne12 * ne13; ++ const uint64_t split_k_size = split_k > 1 ? d_sz * ne12 * ne13 * split_k : 0; ++ if ( ++ (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || ++ (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size) || ++ (split_k > 1 && split_k_size > ctx->device->max_memory_allocation_size)) { ++ GGML_ABORT("Requested preallocation size is too large"); ++ } ++ if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { ++ ctx->prealloc_size_x = x_sz_upd; ++ } ++ if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { ++ ctx->prealloc_size_y = y_sz_upd; ++ } ++ if (split_k > 1 && ctx->prealloc_size_split_k < split_k_size) { ++ ctx->prealloc_size_split_k = split_k_size; ++ } ++ ++ // Request descriptor sets ++ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); ++ if (qx_needs_dequant) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); ++ } ++ if (qy_needs_dequant) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); ++ } ++ if (split_k > 1) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, 1); ++ } ++ return; ++ } ++ ++ vk_buffer d_D = dst_buf_ctx->dev_buffer; ++ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; ++ GGML_ASSERT(d_D != nullptr); ++ GGML_ASSERT(d_D->size >= d_buf_offset + d_sz * ne02 * ne03); ++ vk_buffer d_X; ++ uint64_t x_buf_offset = 0; ++ vk_buffer d_Y; ++ uint64_t y_buf_offset = 0; ++ if (!src0_uma) { ++ d_Qx = src0_buf_ctx->dev_buffer; ++ qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; ++ GGML_ASSERT(d_Qx != nullptr); ++ } ++ if (!src1_uma) { ++ d_Qy = src1_buf_ctx->dev_buffer; ++ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; ++ GGML_ASSERT(d_Qy != nullptr); ++ } ++ if (qx_needs_dequant) { ++ d_X = ctx->prealloc_x; ++ GGML_ASSERT(d_X->size >= x_sz * ne02 * ne03); ++ } else { ++ d_X = d_Qx; ++ x_buf_offset = qx_buf_offset; ++ GGML_ASSERT(qx_sz == x_sz); ++ } ++ if (qy_needs_dequant) { ++ d_Y = ctx->prealloc_y; ++ GGML_ASSERT(d_Y->size >= y_sz * ne02 * ne03); ++ } else { ++ d_Y = d_Qy; ++ y_buf_offset = qy_buf_offset; ++ GGML_ASSERT(qy_sz == y_sz); ++ } ++ ++ if (x_non_contig) { ++ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); ++ } else if (qx_needs_dequant) { ++ const std::vector pc = { (uint32_t)ne01, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)(ggml_nelements(src0)) }; ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz * ne02 * ne03 }, vk_subbuffer{ d_X, 0, x_sz * ne02 * ne03 } }, pc.size() * sizeof(uint32_t), pc.data(), { (uint32_t)(x_ne * ne02 * ne03), 1, 1}); ++ } ++ if (y_non_contig) { ++ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); ++ } ++ ++ uint32_t stride_batch_x = ne00*ne01; ++ uint32_t stride_batch_y = ne10*ne11; ++ ++ if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { ++ stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); ++ } ++ ++ if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { ++ stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); ++ } ++ ++ // compute ++ ggml_vk_matmul( ++ ctx, subctx, pipeline, ++ { d_X, x_buf_offset, x_sz * ne02 * ne03 }, { d_Y, y_buf_offset, y_sz * ne12 * ne13 }, ++ { d_D, d_buf_offset, d_sz * ne12 * ne13 }, { ctx->prealloc_split_k, 0, d_sz * ne12 * ne13 * split_k }, ++ ne01, ne11, ne10, ++ ne10, ne10, ne01, stride_batch_x, stride_batch_y, ne20*ne21, ++ split_k, ne12*ne13, ne02, ne12, r2, r3 ++ ); // NOLINT ++} ++ ++static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ VK_LOG_DEBUG("ggml_vk_mul_mat_vec_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; ++ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; ++ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; ++ std::cerr << "), " << (dryrun ? "dryrun" : "") << "),)"); ++ GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT ++ GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT ++ ++ const uint64_t ne00 = src0->ne[0]; ++ const uint64_t ne01 = src0->ne[1]; ++ const uint64_t ne02 = src0->ne[2]; ++ const uint64_t ne03 = src0->ne[3]; ++ ++ const uint64_t ne10 = src1->ne[0]; ++ const uint64_t ne11 = src1->ne[1]; ++ const uint64_t ne12 = src1->ne[2]; ++ const uint64_t ne13 = src1->ne[3]; ++ ++ const uint64_t ne20 = dst->ne[0]; ++ const uint64_t ne21 = dst->ne[1]; ++ const uint64_t ne22 = dst->ne[2]; ++ const uint64_t ne23 = dst->ne[3]; ++ ++ const uint64_t r2 = ne12 / ne02; ++ const uint64_t r3 = ne13 / ne03; ++ ++ // batch_n indicates that we need to compute a few vector results, and this assumes ++ // ne12 and ne13 are 1. It overloads the batch_strides to hold the row strides. ++ GGML_ASSERT(ne11 == 1 || ne12 * ne13 == 1); ++ bool batch_n = ne11 > 1; ++ ++ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ++ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; ++ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; ++ ++ vk_buffer d_Qx = nullptr; ++ size_t qx_buf_offset = 0; ++ vk_buffer d_Qy = nullptr; ++ size_t qy_buf_offset = 0; ++ ++ bool src0_uma = false; ++ bool src1_uma = false; ++ ++ if (ctx->device->uma) { ++ ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); ++ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); ++ src0_uma = d_Qx != nullptr; ++ src1_uma = d_Qy != nullptr; ++ } ++ ++ const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); ++ const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); ++ ++ const bool f16_f32_kernel = src1->type == GGML_TYPE_F32; ++ ++ const bool qx_needs_dequant = x_non_contig; ++ const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !f16_f32_kernel) || y_non_contig; ++ ++ // Not implemented ++ GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT ++ ++ const uint64_t x_ne = ne01 * ne00; ++ const uint64_t y_ne = ne11 * ne10; ++ const uint64_t d_ne = ne11 * ne01; ++ ++ const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); ++ const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); ++ const uint64_t x_sz = x_non_contig ? ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment) : qx_sz; ++ const uint64_t y_sz = f16_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; ++ const uint64_t d_sz = sizeof(float) * d_ne; ++ ++ vk_pipeline to_fp16_vk_0 = nullptr; ++ vk_pipeline to_fp16_vk_1 = nullptr; ++ if (x_non_contig) { ++ to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, src0->type); ++ } ++ if (y_non_contig) { ++ to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, src1->type); ++ } else { ++ to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); ++ } ++ vk_pipeline dmmv = ggml_vk_get_dequantize_mul_mat_vec(ctx, src0->type, src1->type, ne11); ++ GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT ++ GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT ++ GGML_ASSERT(dmmv != nullptr); ++ ++ if (dryrun) { ++ const uint64_t x_sz_upd = x_sz * ne02 * ne03; ++ const uint64_t y_sz_upd = y_sz * ne12 * ne13; ++ if ( ++ (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || ++ (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { ++ GGML_ABORT("Requested preallocation size is too large"); ++ } ++ if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { ++ ctx->prealloc_size_x = x_sz_upd; ++ } ++ if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { ++ ctx->prealloc_size_y = y_sz_upd; ++ } ++ ++ // Request descriptor sets ++ if (qx_needs_dequant) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); ++ } ++ if (qy_needs_dequant) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); ++ } ++ ggml_pipeline_request_descriptor_sets(ctx->device, dmmv, 1); ++ return; ++ } ++ ++ vk_buffer d_D = dst_buf_ctx->dev_buffer; ++ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; ++ GGML_ASSERT(d_D != nullptr); ++ vk_buffer d_X; ++ uint64_t x_buf_offset = 0; ++ vk_buffer d_Y; ++ uint64_t y_buf_offset = 0; ++ if(!src0_uma) { ++ d_Qx = src0_buf_ctx->dev_buffer; ++ qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; ++ GGML_ASSERT(d_Qx != nullptr); ++ } ++ if(!src1_uma) { ++ d_Qy = src1_buf_ctx->dev_buffer; ++ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; ++ GGML_ASSERT(d_Qy != nullptr); ++ } ++ if (qx_needs_dequant) { ++ d_X = ctx->prealloc_x; ++ } else { ++ d_X = d_Qx; ++ x_buf_offset = qx_buf_offset; ++ GGML_ASSERT(qx_sz == x_sz); ++ } ++ if (qy_needs_dequant) { ++ d_Y = ctx->prealloc_y; ++ } else { ++ d_Y = d_Qy; ++ y_buf_offset = qy_buf_offset; ++ GGML_ASSERT(qy_sz == y_sz); ++ } ++ ++ if (x_non_contig) { ++ GGML_ASSERT(x_sz == ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment)); ++ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); ++ } ++ if (y_non_contig) { ++ GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); ++ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); ++ } ++ ++ // For batch_n, the A matrix is the same for each batch, and B/D use the row stride as the batch stride ++ uint32_t stride_batch_x = batch_n ? 0 : ne00*ne01; ++ uint32_t stride_batch_y = batch_n ? ne10 : (ne10*ne11); ++ uint32_t stride_batch_d = batch_n ? ne20 : (ne20*ne21); ++ ++ if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { ++ stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); ++ } ++ ++ if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { ++ stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); ++ } ++ ++ const uint32_t max_groups_x = ctx->device->properties.limits.maxComputeWorkGroupCount[0]; ++ ++ uint32_t groups_x = ne01; ++ uint32_t groups_z = 1; ++ ++ if (ne01 > max_groups_x) { ++ groups_z = 64; ++ groups_x = CEIL_DIV(groups_x, groups_z); ++ } ++ ++ // compute ++ const vk_mat_vec_push_constants pc = { ++ (uint32_t)ne00, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne01, ++ stride_batch_x, stride_batch_y, stride_batch_d, ++ (uint32_t)ne02, (uint32_t)ne12, (uint32_t)r2, (uint32_t)r3, ++ }; ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, dmmv, ++ { vk_subbuffer{ d_X, x_buf_offset, x_sz * ne02 * ne03 }, vk_subbuffer{ d_Y, y_buf_offset, y_sz * ne12 * ne13 }, vk_subbuffer{ d_D, d_buf_offset, d_sz * ne22 * ne23} }, ++ sizeof(vk_mat_vec_push_constants), &pc, { groups_x, (uint32_t)(ne12 * ne13), groups_z }); ++} ++ ++static void ggml_vk_mul_mat_vec_p021_f16_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ VK_LOG_DEBUG("ggml_vk_mul_mat_p021_f16_f32(" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; ++ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; ++ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; ++ std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); ++ GGML_ASSERT(ggml_is_permuted(src0) && ggml_is_permuted(src1)); ++ GGML_ASSERT(src0->nb[0] <= src0->nb[1] && src0->nb[2] <= src0->nb[3]); // NOLINT ++ GGML_ASSERT(src1->nb[0] <= src1->nb[1] && src1->nb[2] <= src1->nb[3]); // NOLINT ++ GGML_ASSERT(src0->type == GGML_TYPE_F16); ++ GGML_ASSERT(src1->type == GGML_TYPE_F32); ++ ++ const uint64_t ne00 = src0->ne[0]; ++ const uint64_t ne01 = src0->ne[1]; ++ const uint64_t ne02 = src0->ne[2]; ++ // const uint64_t ne03 = src0->ne[3]; ++ ++ const uint64_t ne10 = src1->ne[0]; ++ const uint64_t ne11 = src1->ne[1]; ++ const uint64_t ne12 = src1->ne[2]; ++ // const uint64_t ne13 = src1->ne[3]; ++ ++ GGML_ASSERT(ne11 == 1); ++ ++ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ++ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; ++ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; ++ ++ vk_buffer d_Qy = nullptr; ++ size_t qy_buf_offset = 0; ++ ++ bool src1_uma = false; ++ ++ if (ctx->device->uma) { ++ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); ++ src1_uma = d_Qy != nullptr; ++ } ++ ++ const uint64_t x_ne = ne00 * ne01 * ne02; ++ const uint64_t y_ne = ne10 * ne11 * ne12; ++ const uint64_t d_ne = ne01 * ne11 * ne12; ++ ++ const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); ++ const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); ++ const uint64_t d_sz = sizeof(float) * d_ne; ++ ++ if (dryrun) { ++ // Request descriptor sets ++ ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_mul_mat_vec_p021_f16_f32, 1); ++ return; ++ } ++ ++ vk_buffer d_D = dst_buf_ctx->dev_buffer; ++ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; ++ GGML_ASSERT(d_D != nullptr); ++ vk_buffer d_Qx = src0_buf_ctx->dev_buffer; ++ const uint64_t qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; ++ GGML_ASSERT(d_Qx != nullptr); ++ if (!src1_uma) { ++ d_Qy = src1_buf_ctx->dev_buffer; ++ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; ++ GGML_ASSERT(d_Qx != nullptr); ++ } ++ ++ const uint64_t qy_buffer_offset = (qy_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; ++ const uint64_t qy_shader_offset = qy_buf_offset - qy_buffer_offset; ++ ++ const uint64_t d_buffer_offset = (d_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; ++ const uint64_t d_shader_offset = d_buf_offset - d_buffer_offset; ++ ++ // compute ++ const std::array pc = { (uint32_t)ne00, (uint32_t)ne01, (uint32_t)ne02, (uint32_t)ne12, (uint32_t)(qy_shader_offset / ggml_type_size(src1->type)), (uint32_t)(d_shader_offset / ggml_type_size(dst->type)) }; ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_mul_mat_vec_p021_f16_f32, { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz }, vk_subbuffer{ d_Qy, qy_buffer_offset, qy_sz + qy_shader_offset }, vk_subbuffer{ d_D, d_buffer_offset, d_sz + d_shader_offset } }, 6 * sizeof(uint32_t), &pc, { 1, (uint32_t)ne01, (uint32_t)ne12 }); ++} ++ ++static void ggml_vk_mul_mat_vec_nc_f16_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ VK_LOG_DEBUG("ggml_vk_mul_mat_nc_f16_f32((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; ++ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; ++ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; ++ std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); ++ GGML_ASSERT(!ggml_is_transposed(src0)); ++ GGML_ASSERT(!ggml_is_transposed(src1)); ++ GGML_ASSERT(!ggml_is_permuted(src0)); ++ GGML_ASSERT(src0->type == GGML_TYPE_F16); ++ GGML_ASSERT(src1->type == GGML_TYPE_F32); ++ ++ const uint64_t ne00 = src0->ne[0]; ++ const uint64_t ne01 = src0->ne[1]; ++ const uint64_t ne02 = src0->ne[2]; ++ // const uint64_t ne03 = src0->ne[3]; ++ ++ const uint64_t nb01 = src0->nb[1]; ++ const uint64_t nb02 = src0->nb[2]; ++ ++ // const uint64_t ne10 = src1->ne[0]; ++ const uint64_t ne11 = src1->ne[1]; ++ const uint64_t ne12 = src1->ne[2]; ++ // const uint64_t ne13 = src1->ne[3]; ++ ++ GGML_ASSERT(ne11 == 1); ++ ++ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ++ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; ++ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; ++ ++ vk_buffer d_Qy = nullptr; ++ size_t qy_buf_offset = 0; ++ ++ bool src1_uma = false; ++ ++ if (ctx->device->uma) { ++ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); ++ src1_uma = d_Qy != nullptr; ++ } ++ ++ const uint64_t d_ne = ne01 * ne11 * ne12; ++ ++ const uint32_t row_stride_x = nb01 / sizeof(ggml_fp16_t); ++ const uint32_t channel_stride_x = nb02 / sizeof(ggml_fp16_t); ++ ++ const uint64_t qx_sz = ggml_nbytes(src0); ++ const uint64_t qy_sz = ggml_nbytes(src1); ++ const uint64_t d_sz = sizeof(float) * d_ne; ++ ++ if (dryrun) { ++ // Request descriptor sets ++ ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_mul_mat_vec_nc_f16_f32, 1); ++ return; ++ } ++ ++ vk_buffer d_D = dst_buf_ctx->dev_buffer; ++ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; ++ GGML_ASSERT(d_D != nullptr); ++ vk_buffer d_Qx = src0_buf_ctx->dev_buffer; ++ const uint64_t qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; ++ GGML_ASSERT(d_Qx != nullptr); ++ if (!src1_uma) { ++ d_Qy = src1_buf_ctx->dev_buffer; ++ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; ++ GGML_ASSERT(d_Qx != nullptr); ++ } ++ ++ const uint64_t qy_buffer_offset = (qy_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; ++ const uint64_t qy_shader_offset = qy_buf_offset - qy_buffer_offset; ++ ++ const uint64_t d_buffer_offset = (d_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; ++ const uint64_t d_shader_offset = d_buf_offset - d_buffer_offset; ++ ++ // compute ++ const std::array pc = { (uint32_t)ne00, (uint32_t)ne01, row_stride_x, channel_stride_x, (uint32_t)(ne12 / ne02), (uint32_t)(qy_shader_offset / ggml_type_size(src1->type)), (uint32_t)(d_shader_offset / ggml_type_size(dst->type)) }; ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_mul_mat_vec_nc_f16_f32, ++ { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz }, vk_subbuffer{ d_Qy, qy_buffer_offset, qy_sz + qy_shader_offset }, vk_subbuffer{ d_D, d_buffer_offset, d_sz + d_shader_offset } }, 7 * sizeof(uint32_t), &pc, { 1, (uint32_t)ne01, (uint32_t)ne12 }); ++} ++ ++static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ VK_LOG_DEBUG("ggml_vk_mul_mat(" << src0 << ", " << src1 << ", " << dst << ")"); ++ if (src0->type == GGML_TYPE_F16 && ggml_is_permuted(src0) && ggml_is_permuted(src1) && dst->ne[1] == 1 && ++ // detect 0213 permutation, and batch size of 1 ++ src0->nb[0] <= src0->nb[2] && ++ src0->nb[2] <= src0->nb[1] && ++ src0->nb[1] <= src0->nb[3] && ++ src1->nb[0] <= src1->nb[2] && ++ src1->nb[2] <= src1->nb[1] && ++ src1->nb[1] <= src1->nb[3] && ++ src0->ne[3] == 1 && ++ src1->ne[3] == 1) { ++ ggml_vk_mul_mat_vec_p021_f16_f32(ctx, subctx, src0, src1, dst, dryrun); ++ } else if (src0->type == GGML_TYPE_F16 && !ggml_is_contiguous(src0) && !ggml_is_transposed(src1) && dst->ne[1] == 1 && ++ !ggml_is_permuted(src0) && !ggml_is_permuted(src1)) { ++ ggml_vk_mul_mat_vec_nc_f16_f32(ctx, subctx, src0, src1, dst, dryrun); ++ // mul_mat_vec supports batching ne12*ne13 when ne11==1, or treating ne11 as the batch size (up to four) ++ // when ne12 and ne13 are one. ++ } else if ((dst->ne[1] == 1 || (dst->ne[1] <= mul_mat_vec_max_cols && src1->ne[2] * src1->ne[3] == 1)) && ++ (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type))) { ++ ggml_vk_mul_mat_vec_q_f16(ctx, subctx, src0, src1, dst, dryrun); ++ } else { ++ ggml_vk_mul_mat_q_f16(ctx, subctx, src0, src1, dst, dryrun); ++ } ++} ++ ++static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, bool dryrun = false) { ++ VK_LOG_DEBUG("ggml_vk_mul_mat_id_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; ++ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; ++ std::cerr << "), (" << ids << ", name=" << ids->name << ", type=" << ids->type << ", ne0=" << ids->ne[0] << ", ne1=" << ids->ne[1] << ", ne2=" << ids->ne[2] << ", ne3=" << ids->ne[3] << ", nb0=" << ids->nb[0] << ", nb1=" << ids->nb[1] << ", nb2=" << ids->nb[2] << ", nb3=" << ids->nb[3]; ++ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3] << "),)"); ++ GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT ++ GGML_ASSERT(ids->type == GGML_TYPE_I32); ++ ++ const uint64_t ne00 = src0->ne[0]; ++ const uint64_t ne01 = src0->ne[1]; ++ const uint64_t ne02 = src0->ne[2]; ++ const uint64_t ne03 = src0->ne[3]; ++ ++ const uint64_t ne10 = src1->ne[0]; ++ const uint64_t ne11 = src1->ne[1]; ++ const uint64_t ne12 = src1->ne[2]; ++ const uint64_t ne13 = src1->ne[3]; ++ ++ const uint64_t nei0 = ids->ne[0]; ++ const uint64_t nei1 = ids->ne[1]; ++ GGML_ASSERT(nei0 * nei1 <= 3072); ++ ++ const uint32_t nbi1 = ids->nb[1]; ++ const uint32_t nbi2 = ids->nb[2]; ++ ++ const uint64_t ne20 = dst->ne[0]; ++ const uint64_t ne21 = dst->ne[1]; ++ const uint64_t ne22 = dst->ne[2]; ++ const uint64_t ne23 = dst->ne[3]; ++ ++ const uint64_t n_as = ne02; ++ ++ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ++ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; ++ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; ++ ggml_backend_vk_buffer_context * ids_buf_ctx = (ggml_backend_vk_buffer_context *)ids->buffer->context; ++ ++ vk_buffer d_Qx = nullptr; ++ size_t qx_buf_offset = 0; ++ vk_buffer d_Qy = nullptr; ++ size_t qy_buf_offset = 0; ++ vk_buffer d_ids = nullptr; ++ size_t ids_buf_offset = 0; ++ ++ bool src0_uma = false; ++ bool src1_uma = false; ++ bool ids_uma = false; ++ ++ if (ctx->device->uma) { ++ ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); ++ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); ++ ggml_vk_host_get(ctx->device, ids->data, d_ids, ids_buf_offset); ++ src0_uma = d_Qx != nullptr; ++ src1_uma = d_Qy != nullptr; ++ ids_uma = d_ids != nullptr; ++ } ++ ++ const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); ++ const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); ++ ++ const bool y_f32_kernel = src1->type == GGML_TYPE_F32 && !y_non_contig; ++ ++ vk_matmul_pipeline mmp = ggml_vk_get_mul_mat_mat_id_pipeline(ctx, src0->type, y_non_contig ? GGML_TYPE_F16 : src1->type, (ggml_prec)dst->op_params[0]); ++ ++ const bool qx_needs_dequant = mmp == nullptr || x_non_contig; ++ const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !y_f32_kernel) || y_non_contig; ++ ++ if (qx_needs_dequant) { ++ GGML_ABORT("fatal error"); ++ } ++ ++ // Not implemented ++ GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT ++ ++ const uint64_t x_ne = ne01 * ne00; ++ const uint64_t y_ne = ne11 * ne10; ++ const uint64_t d_ne = ne21 * ne20; ++ ++ const uint32_t kpad = ggml_vk_align_size(ne10, ggml_vk_guess_matmul_id_pipeline_align(ctx, mmp, ne01, nei1)); ++ const bool aligned = ne10 == kpad && ne01 > 8 && nei1 > 8; ++ ++ vk_pipeline pipeline = ggml_vk_guess_matmul_id_pipeline(ctx, mmp, ne01, nei1, aligned); ++ ++ const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type); ++ const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); ++ const uint64_t x_sz = !qx_needs_dequant ? qx_sz : sizeof(ggml_fp16_t) * x_ne; ++ const uint64_t y_sz = y_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; ++ const uint64_t ids_sz = nbi2; ++ const uint64_t d_sz = sizeof(float) * d_ne; ++ ++ vk_pipeline to_fp16_vk_0 = nullptr; ++ vk_pipeline to_fp16_vk_1 = nullptr; ++ ++ if (x_non_contig) { ++ to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, GGML_TYPE_F16); ++ } else { ++ to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, src0->type); ++ } ++ if (y_non_contig) { ++ to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, GGML_TYPE_F16); ++ } else { ++ to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); ++ } ++ GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT ++ GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT ++ ++ if (dryrun) { ++ const uint64_t x_sz_upd = x_sz * ne02 * ne03; ++ const uint64_t y_sz_upd = y_sz * ne12 * ne13; ++ if ( ++ (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || ++ (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { ++ GGML_ABORT("Requested preallocation size is too large"); ++ } ++ if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { ++ ctx->prealloc_size_x = x_sz_upd; ++ } ++ if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { ++ ctx->prealloc_size_y = y_sz_upd; ++ } ++ ++ // Request descriptor sets ++ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); ++ if (qx_needs_dequant) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); ++ } ++ if (qy_needs_dequant) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); ++ } ++ return; ++ } ++ ++ vk_buffer d_D = dst_buf_ctx->dev_buffer; ++ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; ++ GGML_ASSERT(d_D != nullptr); ++ vk_buffer d_X; ++ uint64_t x_buf_offset = 0; ++ vk_buffer d_Y; ++ uint64_t y_buf_offset = 0; ++ if (!src0_uma) { ++ d_Qx = src0_buf_ctx->dev_buffer; ++ qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; ++ GGML_ASSERT(d_Qx != nullptr); ++ } ++ if (!src1_uma) { ++ d_Qy = src1_buf_ctx->dev_buffer; ++ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; ++ GGML_ASSERT(d_Qy != nullptr); ++ } ++ if (!ids_uma) { ++ d_ids = ids_buf_ctx->dev_buffer; ++ ids_buf_offset = vk_tensor_offset(ids) + ids->view_offs; ++ GGML_ASSERT(d_ids != nullptr); ++ } ++ if (qx_needs_dequant) { ++ d_X = ctx->prealloc_x; ++ GGML_ASSERT(d_X->size >= x_sz * ne02 * ne03); ++ } else { ++ d_X = d_Qx; ++ x_buf_offset = qx_buf_offset; ++ GGML_ASSERT(qx_sz == x_sz); ++ } ++ if (qy_needs_dequant) { ++ d_Y = ctx->prealloc_y; ++ GGML_ASSERT(d_Y->size >= y_sz * ne02 * ne03); ++ } else { ++ d_Y = d_Qy; ++ y_buf_offset = qy_buf_offset; ++ GGML_ASSERT(qy_sz == y_sz); ++ } ++ ++ if (x_non_contig) { ++ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); ++ } else if (qx_needs_dequant) { ++ const std::vector pc = { (uint32_t)ne01, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)(ggml_nelements(src0)) }; ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, ++ { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz * ne02 * ne03 }, vk_subbuffer{ d_X, 0, x_sz * ne02 * ne03 } }, pc.size() * sizeof(uint32_t), pc.data(), { (uint32_t)(x_ne * ne02 * ne03), 1, 1}); ++ } ++ if (y_non_contig) { ++ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); ++ } ++ ++ uint32_t stride_batch_x = ne00*ne01; ++ uint32_t stride_batch_y = ne10*ne11; ++ ++ if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { ++ stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); ++ } ++ ++ if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { ++ stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); ++ } ++ ++ // compute ++ ggml_vk_matmul_id( ++ ctx, subctx, pipeline, ++ { d_X, x_buf_offset, x_sz * ne02 * ne03 }, { d_Y, y_buf_offset, y_sz * ne12 * ne13 }, ++ { d_D, d_buf_offset, d_sz * ne22 * ne23 }, { d_ids, ids_buf_offset, ids_sz }, ++ ne01, ne21, ne10, ne10, ne10, ne01, ++ stride_batch_x, stride_batch_y, ne20*ne21, ++ n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11 ++ ); // NOLINT ++} ++ ++static void ggml_vk_mul_mat_vec_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, bool dryrun = false) { ++ VK_LOG_DEBUG("ggml_vk_mul_mat_vec_id_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; ++ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; ++ std::cerr << "), (" << ids << ", name=" << ids->name << ", type=" << ids->type << ", ne0=" << ids->ne[0] << ", ne1=" << ids->ne[1] << ", ne2=" << ids->ne[2] << ", ne3=" << ids->ne[3] << ", nb0=" << ids->nb[0] << ", nb1=" << ids->nb[1] << ", nb2=" << ids->nb[2] << ", nb3=" << ids->nb[3]; ++ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; ++ std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); ++ GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT ++ GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT ++ GGML_ASSERT(ids->type == GGML_TYPE_I32); ++ ++ const uint64_t ne00 = src0->ne[0]; ++ const uint64_t ne01 = src0->ne[1]; ++ const uint64_t ne02 = src0->ne[2]; ++ const uint64_t ne03 = src0->ne[3]; ++ ++ const uint64_t ne10 = src1->ne[0]; ++ const uint64_t ne11 = src1->ne[1]; ++ const uint64_t ne12 = src1->ne[2]; ++ const uint64_t ne13 = src1->ne[3]; ++ ++ const uint64_t nei0 = ids->ne[0]; ++ const uint64_t nei1 = ids->ne[1]; ++ ++ const uint64_t nbi2 = ids->nb[2]; ++ ++ GGML_ASSERT(nei1 == 1); ++ ++ const uint64_t ne20 = dst->ne[0]; ++ const uint64_t ne21 = dst->ne[1]; ++ const uint64_t ne22 = dst->ne[2]; ++ const uint64_t ne23 = dst->ne[3]; ++ ++ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ++ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; ++ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; ++ ggml_backend_vk_buffer_context * ids_buf_ctx = (ggml_backend_vk_buffer_context *)ids->buffer->context; ++ ++ vk_buffer d_Qx = nullptr; ++ size_t qx_buf_offset = 0; ++ vk_buffer d_Qy = nullptr; ++ size_t qy_buf_offset = 0; ++ vk_buffer d_ids = nullptr; ++ size_t ids_buf_offset = 0; ++ ++ bool src0_uma = false; ++ bool src1_uma = false; ++ bool ids_uma = false; ++ ++ if (ctx->device->uma) { ++ ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); ++ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); ++ ggml_vk_host_get(ctx->device, ids->data, d_ids, ids_buf_offset); ++ src0_uma = d_Qx != nullptr; ++ src1_uma = d_Qy != nullptr; ++ ids_uma = d_ids != nullptr; ++ } ++ ++ const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); ++ const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); ++ ++ const bool f16_f32_kernel = src1->type == GGML_TYPE_F32; ++ ++ const bool qx_needs_dequant = x_non_contig; ++ const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !f16_f32_kernel) || y_non_contig; ++ ++ // Not implemented ++ GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT ++ ++ const uint64_t x_ne = ne01 * ne00; ++ const uint64_t y_ne = ne11 * ne10; ++ const uint64_t d_ne = ne21 * ne20; ++ ++ const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); ++ const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); ++ const uint64_t x_sz = x_non_contig ? ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment) : qx_sz; ++ const uint64_t y_sz = f16_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; ++ const uint64_t ids_sz = nbi2; ++ const uint64_t d_sz = sizeof(float) * d_ne; ++ ++ vk_pipeline to_fp16_vk_0 = nullptr; ++ vk_pipeline to_fp16_vk_1 = nullptr; ++ if (x_non_contig) { ++ to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, src0->type); ++ } ++ if (y_non_contig) { ++ to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, src1->type); ++ } else { ++ to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); ++ } ++ vk_pipeline dmmv = ggml_vk_get_dequantize_mul_mat_vec_id(ctx, src0->type, src1->type); ++ GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT ++ GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT ++ GGML_ASSERT(dmmv != nullptr); ++ ++ if (dryrun) { ++ const uint64_t x_sz_upd = x_sz * ne02 * ne03; ++ const uint64_t y_sz_upd = y_sz * ne12 * ne13; ++ if ( ++ (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || ++ (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { ++ GGML_ABORT("Requested preallocation size is too large"); ++ } ++ if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { ++ ctx->prealloc_size_x = x_sz_upd; ++ } ++ if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { ++ ctx->prealloc_size_y = y_sz_upd; ++ } ++ ++ // Request descriptor sets ++ if (qx_needs_dequant) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); ++ } ++ if (qy_needs_dequant) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); ++ } ++ ggml_pipeline_request_descriptor_sets(ctx->device, dmmv, 1); ++ return; ++ } ++ ++ vk_buffer d_D = dst_buf_ctx->dev_buffer; ++ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; ++ GGML_ASSERT(d_D != nullptr); ++ vk_buffer d_X; ++ uint64_t x_buf_offset = 0; ++ vk_buffer d_Y; ++ uint64_t y_buf_offset = 0; ++ if(!src0_uma) { ++ d_Qx = src0_buf_ctx->dev_buffer; ++ qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; ++ GGML_ASSERT(d_Qx != nullptr); ++ } ++ if(!src1_uma) { ++ d_Qy = src1_buf_ctx->dev_buffer; ++ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; ++ GGML_ASSERT(d_Qy != nullptr); ++ } ++ if(!ids_uma) { ++ d_ids = ids_buf_ctx->dev_buffer; ++ ids_buf_offset = vk_tensor_offset(ids) + ids->view_offs; ++ GGML_ASSERT(d_ids != nullptr); ++ } ++ if (qx_needs_dequant) { ++ d_X = ctx->prealloc_x; ++ } else { ++ d_X = d_Qx; ++ x_buf_offset = qx_buf_offset; ++ GGML_ASSERT(qx_sz == x_sz); ++ } ++ if (qy_needs_dequant) { ++ d_Y = ctx->prealloc_y; ++ } else { ++ d_Y = d_Qy; ++ y_buf_offset = qy_buf_offset; ++ GGML_ASSERT(qy_sz == y_sz); ++ } ++ ++ if (x_non_contig) { ++ GGML_ASSERT(x_sz == ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment)); ++ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); ++ } ++ if (y_non_contig) { ++ GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); ++ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); ++ } ++ ++ uint32_t stride_batch_y = ne10*ne11; ++ ++ if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { ++ stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); ++ } ++ ++ const uint32_t max_groups_x = ctx->device->properties.limits.maxComputeWorkGroupCount[0]; ++ ++ uint32_t groups_x = ne01; ++ uint32_t groups_z = 1; ++ ++ if (ne01 > max_groups_x) { ++ groups_z = 64; ++ groups_x = CEIL_DIV(groups_x, groups_z); ++ } ++ ++ // compute ++ const vk_mat_vec_id_push_constants pc = { ++ (uint32_t)ne00, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne01, ++ (uint32_t)x_ne, stride_batch_y, (uint32_t)(ne20*ne21), ++ (uint32_t)nei0, (uint32_t)ne11, ++ }; ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, dmmv, ++ { vk_subbuffer{ d_X, x_buf_offset, x_sz * ne02 * ne03 }, ++ vk_subbuffer{ d_Y, y_buf_offset, y_sz * ne12 * ne13 }, vk_subbuffer{ d_D, d_buf_offset, d_sz * ne22 * ne23}, vk_subbuffer{ d_ids, ids_buf_offset, ids_sz } }, ++ sizeof(vk_mat_vec_id_push_constants), &pc, { groups_x, (uint32_t)nei0, groups_z }); ++} ++ ++static void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, bool dryrun = false) { ++ VK_LOG_DEBUG("ggml_vk_mul_mat_id(" << src0 << ", " << src1 << ", " << src2 << ", " << dst << ")"); ++ if (src2->ne[1] == 1 && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type))) { ++ ggml_vk_mul_mat_vec_id_q_f16(ctx, subctx, src0, src1, src2, dst, dryrun); ++ } else { ++ ggml_vk_mul_mat_id_q_f16(ctx, subctx, src0, src1, src2, dst, dryrun); ++ } ++} ++ ++static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, ggml_tensor * dst, bool dryrun = false) { ++ VK_LOG_DEBUG("ggml_vk_flash_attn((" << q << ", name=" << q->name << ", type=" << q->type << ", ne0=" << q->ne[0] << ", ne1=" << q->ne[1] << ", ne2=" << q->ne[2] << ", ne3=" << q->ne[3] << ", nb0=" << q->nb[0] << ", nb1=" << q->nb[1] << ", nb2=" << q->nb[2] << ", nb3=" << q->nb[3]; ++ std::cerr << "), (" << k << ", name=" << k->name << ", type=" << k->type << ", ne0=" << k->ne[0] << ", ne1=" << k->ne[1] << ", ne2=" << k->ne[2] << ", ne3=" << k->ne[3] << ", nb0=" << k->nb[0] << ", nb1=" << k->nb[1] << ", nb2=" << k->nb[2] << ", nb3=" << k->nb[3]; ++ std::cerr << "), (" << v << ", name=" << v->name << ", type=" << v->type << ", ne0=" << v->ne[0] << ", ne1=" << v->ne[1] << ", ne2=" << v->ne[2] << ", ne3=" << v->ne[3] << ", nb0=" << v->nb[0] << ", nb1=" << v->nb[1] << ", nb2=" << v->nb[2] << ", nb3=" << v->nb[3]; ++ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; ++ std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); ++ ++ GGML_TENSOR_LOCALS(int64_t, neq, q, ne) ++ GGML_TENSOR_LOCALS(size_t, nbq, q, nb) ++ GGML_TENSOR_LOCALS(int64_t, nek, k, ne) ++ GGML_TENSOR_LOCALS(size_t, nbk, k, nb) ++ GGML_TENSOR_LOCALS(int64_t, nev, v, ne) ++ GGML_TENSOR_LOCALS(size_t, nbv, v, nb) ++ GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) ++ GGML_TENSOR_LOCALS(size_t, nb, dst, nb) ++ ++ const uint32_t nem1 = mask ? mask->ne[1] : 0; ++ const uint32_t nbm1 = mask ? mask->nb[1] : 0; ++ ++ const uint32_t D = neq0; ++ const uint32_t N = neq1; ++ const uint32_t KV = nek1; ++ ++ GGML_ASSERT(ne0 == D); ++ GGML_ASSERT(ne2 == N); ++ ++ // input tensor rows must be contiguous ++ GGML_ASSERT(nbq0 == ggml_type_size(q->type)); ++ GGML_ASSERT(nbk0 == ggml_type_size(k->type)); ++ GGML_ASSERT(nbv0 == ggml_type_size(v->type)); ++ ++ GGML_ASSERT(neq0 == D); ++ GGML_ASSERT(nek0 == D); ++ GGML_ASSERT(nev0 == D); ++ ++ GGML_ASSERT(neq1 == N); ++ GGML_ASSERT(nev0 == D); ++ ++ GGML_ASSERT(nev1 == nek1); ++ ++ // dst cannot be transposed or permuted ++ GGML_ASSERT(nb0 == sizeof(float)); ++ GGML_ASSERT(nb0 <= nb1); ++ GGML_ASSERT(nb1 <= nb2); ++ GGML_ASSERT(nb2 <= nb3); ++ ++ assert(dst->type == GGML_TYPE_F32); ++ assert(q->type == GGML_TYPE_F32); ++ assert(k->type == v->type); ++ ++ vk_pipeline *pipelines; ++ // XXX TODO other backends may be changing accumulator precision to default to f32 soon ++ bool f32acc = dst->op_params[3] == GGML_PREC_F32; ++ bool small_rows = N <= flash_attention_num_small_rows; ++ switch (D) { ++ case 64: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D64[k->type][f32acc][small_rows][0]; break; ++ case 80: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D80[k->type][f32acc][small_rows][0]; break; ++ case 96: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D96[k->type][f32acc][small_rows][0]; break; ++ case 112: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D112[k->type][f32acc][small_rows][0]; break; ++ case 128: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D128[k->type][f32acc][small_rows][0]; break; ++ case 256: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D256[k->type][f32acc][small_rows][0]; break; ++ default: ++ assert(!"unsupported D value"); ++ return; ++ } ++ assert(pipelines); ++ ++ bool aligned = (KV % pipelines[1]->align) == 0; ++ vk_pipeline pipeline = pipelines[aligned]; ++ assert(pipeline); ++ ++ if (dryrun) { ++ // Request descriptor sets ++ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); ++ return; ++ } ++ ++ float scale = 1.0f; ++ float max_bias = 0.0f; ++ float logit_softcap = 0.0f; ++ ++ memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); ++ memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); ++ memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); ++ ++ if (logit_softcap != 0) { ++ scale /= logit_softcap; ++ } ++ ++ const uint32_t n_head_kv = neq2; ++ const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head_kv)); ++ const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); ++ const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); ++ ++ ggml_vk_sync_buffers(subctx); ++ ++ vk_buffer d_Q = nullptr, d_K = nullptr, d_V = nullptr, d_D = nullptr, d_M = nullptr; ++ size_t q_buf_offset = 0, k_buf_offset = 0, v_buf_offset = 0, d_buf_offset = 0, m_buf_offset = 0; ++ ++ bool Q_uma = false, K_uma = false, V_uma = false, D_uma = false, M_uma = false; ++ ++ if (ctx->device->uma) { ++ ggml_vk_host_get(ctx->device, q->data, d_Q, q_buf_offset); ++ ggml_vk_host_get(ctx->device, k->data, d_K, q_buf_offset); ++ ggml_vk_host_get(ctx->device, v->data, d_V, q_buf_offset); ++ ggml_vk_host_get(ctx->device, dst->data, d_D, q_buf_offset); ++ Q_uma = d_Q != nullptr; ++ K_uma = d_K != nullptr; ++ V_uma = d_V != nullptr; ++ D_uma = d_D != nullptr; ++ if (mask) { ++ ggml_vk_host_get(ctx->device, mask->data, d_M, q_buf_offset); ++ M_uma = d_M != nullptr; ++ } ++ } ++ ++ ++ ggml_backend_vk_buffer_context * d_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ++ ggml_backend_vk_buffer_context * q_buf_ctx = (ggml_backend_vk_buffer_context *)q->buffer->context; ++ ggml_backend_vk_buffer_context * k_buf_ctx = (ggml_backend_vk_buffer_context *)k->buffer->context; ++ ggml_backend_vk_buffer_context * v_buf_ctx = (ggml_backend_vk_buffer_context *)v->buffer->context; ++ ++ if (!Q_uma) { ++ d_Q = q_buf_ctx->dev_buffer; ++ q_buf_offset = vk_tensor_offset(q) + q->view_offs; ++ } ++ if (!K_uma) { ++ d_K = k_buf_ctx->dev_buffer; ++ k_buf_offset = vk_tensor_offset(k) + k->view_offs; ++ } ++ if (!V_uma) { ++ d_V = v_buf_ctx->dev_buffer; ++ v_buf_offset = vk_tensor_offset(v) + v->view_offs; ++ } ++ if (!D_uma) { ++ d_D = d_buf_ctx->dev_buffer; ++ d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; ++ } ++ ++ if (!M_uma) { ++ d_M = d_Q; ++ m_buf_offset = q_buf_offset; ++ if (mask) { ++ ggml_backend_vk_buffer_context * m_buf_ctx = (ggml_backend_vk_buffer_context*)mask->buffer->context; ++ d_M = m_buf_ctx->dev_buffer; ++ m_buf_offset = vk_tensor_offset(mask) + mask->view_offs; ++ } ++ } ++ ++ const vk_flash_attn_push_constants pc = { N, KV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3, (uint32_t)neq2, (uint32_t)neq3, (uint32_t)nek2, (uint32_t)nek3, (uint32_t)nev2, (uint32_t)nev3, nem1, (uint32_t)nbq2, (uint32_t)nbq3, (uint32_t)nbk2, (uint32_t)nbk3, (uint32_t)nbv2, (uint32_t)nbv3, nbm1, scale, max_bias, logit_softcap, mask != nullptr, n_head_log2, m0, m1 }; ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, ++ { ++ vk_subbuffer{d_Q, q_buf_offset, VK_WHOLE_SIZE}, ++ vk_subbuffer{d_K, k_buf_offset, VK_WHOLE_SIZE}, ++ vk_subbuffer{d_V, v_buf_offset, VK_WHOLE_SIZE}, ++ vk_subbuffer{d_M, m_buf_offset, VK_WHOLE_SIZE}, ++ vk_subbuffer{d_D, d_buf_offset, VK_WHOLE_SIZE}, ++ }, ++ sizeof(vk_flash_attn_push_constants), &pc, { (uint32_t)neq1, (uint32_t)neq2, (uint32_t)neq3 }); ++} ++ ++static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, ggml_op op) { ++ switch (op) { ++ case GGML_OP_GET_ROWS: ++ GGML_ASSERT(src1->type == GGML_TYPE_I32); ++ if (dst->type == GGML_TYPE_F16) { ++ return ctx->device->pipeline_get_rows[src0->type]; ++ } ++ if (dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_get_rows_f32[src0->type]; ++ } ++ return nullptr; ++ case GGML_OP_ACC: ++ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_acc_f32; ++ } ++ return nullptr; ++ case GGML_OP_ADD: ++ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_add_f32_norepeat : ctx->device->pipeline_add_f32; ++ } ++ if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) { ++ return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_add_f16_f32_f16_norepeat : ctx->device->pipeline_add_f16_f32_f16; ++ } ++ return nullptr; ++ case GGML_OP_MUL: ++ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_mul_f32_norepeat : ctx->device->pipeline_mul_f32; ++ } ++ return nullptr; ++ case GGML_OP_DIV: ++ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_div_f32_norepeat : ctx->device->pipeline_div_f32; ++ } ++ return nullptr; ++ case GGML_OP_CONCAT: ++ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_concat_f32; ++ } ++ if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { ++ return ctx->device->pipeline_concat_f16; ++ } ++ if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_I32 && dst->type == GGML_TYPE_I32) { ++ return ctx->device->pipeline_concat_i32; ++ } ++ return nullptr; ++ case GGML_OP_UPSCALE: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_upscale_f32; ++ } ++ return nullptr; ++ case GGML_OP_SCALE: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_scale_f32; ++ } ++ return nullptr; ++ case GGML_OP_SQR: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_sqr_f32; ++ } ++ return nullptr; ++ case GGML_OP_SIN: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_sin_f32; ++ } ++ return nullptr; ++ case GGML_OP_COS: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_cos_f32; ++ } ++ return nullptr; ++ case GGML_OP_CLAMP: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_clamp_f32; ++ } ++ return nullptr; ++ case GGML_OP_PAD: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_pad_f32; ++ } ++ return nullptr; ++ case GGML_OP_REPEAT: ++ if (ggml_type_size(src0->type) == sizeof(float) && ggml_type_size(dst->type) == sizeof(float)) { ++ return ctx->device->pipeline_repeat_f32; ++ } ++ return nullptr; ++ case GGML_OP_CPY: ++ case GGML_OP_CONT: ++ case GGML_OP_DUP: ++ return ggml_vk_get_cpy_pipeline(ctx, src0, dst, dst->type); ++ case GGML_OP_NORM: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_norm_f32; ++ } ++ return nullptr; ++ case GGML_OP_GROUP_NORM: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_group_norm_f32; ++ } ++ return nullptr; ++ case GGML_OP_RMS_NORM: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_rms_norm_f32; ++ } ++ return nullptr; ++ case GGML_OP_UNARY: ++ switch (ggml_get_unary_op(dst)) { ++ case GGML_UNARY_OP_SILU: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_silu_f32; ++ } ++ break; ++ case GGML_UNARY_OP_GELU: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_gelu_f32; ++ } ++ break; ++ case GGML_UNARY_OP_GELU_QUICK: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_gelu_quick_f32; ++ } ++ break; ++ case GGML_UNARY_OP_RELU: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_relu_f32; ++ } ++ break; ++ case GGML_UNARY_OP_TANH: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_tanh_f32; ++ } ++ break; ++ default: ++ break; ++ } ++ return nullptr; ++ case GGML_OP_DIAG_MASK_INF: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_diag_mask_inf_f32; ++ } ++ return nullptr; ++ case GGML_OP_SOFT_MAX: ++ GGML_ASSERT(!src1 || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); ++ ++ if (src0->type == GGML_TYPE_F32 && (src1 == nullptr || src1->type == GGML_TYPE_F32) && dst->type == GGML_TYPE_F32) { ++ return src0->ne[0] > 1024 ? ctx->device->pipeline_soft_max_f32_wg512 : ctx->device->pipeline_soft_max_f32; ++ } ++ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) { ++ return src0->ne[0] > 1024 ? ctx->device->pipeline_soft_max_f32_f16_wg512 : ctx->device->pipeline_soft_max_f32_f16; ++ } ++ return nullptr; ++ case GGML_OP_ROPE: ++ { ++ const int mode = ((const int32_t *) dst->op_params)[2]; ++ const bool is_neox = mode & GGML_ROPE_TYPE_NEOX; ++ ++ if (is_neox) { ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_rope_neox_f32; ++ } ++ if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { ++ return ctx->device->pipeline_rope_neox_f16; ++ } ++ } else { ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_rope_norm_f32; ++ } ++ if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { ++ return ctx->device->pipeline_rope_norm_f16; ++ } ++ } ++ return nullptr; ++ } ++ case GGML_OP_ARGSORT: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_I32) { ++ return ctx->device->pipeline_argsort_f32; ++ } ++ return nullptr; ++ case GGML_OP_SUM_ROWS: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_sum_rows_f32; ++ } ++ return nullptr; ++ case GGML_OP_IM2COL: ++ if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_im2col_f32; ++ } ++ if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) { ++ return ctx->device->pipeline_im2col_f32_f16; ++ } ++ return nullptr; ++ case GGML_OP_TIMESTEP_EMBEDDING: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_timestep_embedding_f32; ++ } ++ return nullptr; ++ case GGML_OP_POOL_2D: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_pool2d_f32; ++ } ++ return nullptr; ++ case GGML_OP_RWKV_WKV6: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_rwkv_wkv6_f32; ++ } ++ return nullptr; ++ case GGML_OP_LEAKY_RELU: ++ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ++ return ctx->device->pipeline_leaky_relu_f32; ++ } ++ return nullptr; ++ default: ++ return nullptr; ++ } ++ ++ GGML_UNUSED(src2); ++} ++ ++static bool ggml_vk_op_supports_incontiguous(ggml_op op) { ++ switch (op) { ++ case GGML_OP_CPY: ++ case GGML_OP_GET_ROWS: ++ case GGML_OP_ADD: ++ case GGML_OP_MUL: ++ case GGML_OP_DIV: ++ case GGML_OP_CONCAT: ++ case GGML_OP_UPSCALE: ++ case GGML_OP_SQR: ++ case GGML_OP_SIN: ++ case GGML_OP_COS: ++ case GGML_OP_CLAMP: ++ case GGML_OP_PAD: ++ case GGML_OP_REPEAT: ++ return true; ++ default: ++ return false; ++ } ++} ++ ++static uint32_t get_misalign_bytes(ggml_backend_vk_context * ctx, const ggml_tensor * t) ++{ ++ return ((vk_tensor_offset(t) + t->view_offs) & (ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1));; ++} ++ ++template void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, T &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { ++ GGML_UNUSED(p); ++ GGML_UNUSED(src0); ++ GGML_UNUSED(src1); ++ GGML_UNUSED(src2); ++ GGML_UNUSED(dst); ++ static_assert(!std::is_const::value, "unexpected type"); ++ GGML_ASSERT(!src0 || get_misalign_bytes(ctx, src0) == 0); ++ GGML_ASSERT(!src1 || get_misalign_bytes(ctx, src1) == 0); ++ GGML_ASSERT(!src2 || get_misalign_bytes(ctx, src2) == 0); ++ GGML_ASSERT(!dst || get_misalign_bytes(ctx, dst) == 0); ++} ++ ++template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_unary_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { ++ const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); ++ const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); ++ ++ p.misalign_offsets = (a_offset << 16) | d_offset; ++ ++ GGML_UNUSED(src1); ++ GGML_UNUSED(src2); ++} ++ ++template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_binary_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { ++ const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); ++ const uint32_t b_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); ++ const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); ++ ++ GGML_ASSERT(dst->op != GGML_OP_GET_ROWS || (a_offset == 0 && b_offset == 0 && d_offset == 0)); ++ ++ p.misalign_offsets = (a_offset << 16) | (b_offset << 8) | d_offset; ++ ++ GGML_UNUSED(src2); ++} ++ ++template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_upscale_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { ++ const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); ++ const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); ++ ++ p.a_offset = a_offset; ++ p.d_offset = d_offset; ++ ++ GGML_UNUSED(src1); ++ GGML_UNUSED(src2); ++} ++ ++template ++static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, ggml_op op, PC&& pc, bool dryrun = false) { ++ VK_LOG_DEBUG("ggml_vk_op_f32((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; ++ if (src1 != nullptr) { ++ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; ++ } ++ if (src2 != nullptr) { ++ std::cerr << "), (" << src2 << ", name=" << src2->name << ", type=" << src2->type << ", ne0=" << src2->ne[0] << ", ne1=" << src2->ne[1] << ", ne2=" << src2->ne[2] << ", ne3=" << src2->ne[3] << ", nb0=" << src2->nb[0] << ", nb1=" << src2->nb[1] << ", nb2=" << src2->nb[2] << ", nb3=" << src2->nb[3]; ++ } ++ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; ++ std::cerr << "), " << ggml_op_name(op) << ", " << (dryrun ? "dryrun" : "") << ")"); ++ GGML_ASSERT(op == GGML_OP_GET_ROWS || (!ggml_is_quantized(src0->type) && (src1 == nullptr || !ggml_is_quantized(src1->type)))); // NOLINT ++ GGML_ASSERT(ggml_vk_op_supports_incontiguous(op) || ggml_vk_dim01_contiguous(src0)); // NOLINT ++ GGML_ASSERT(dst->buffer != nullptr); ++ const uint64_t ne00 = src0->ne[0]; ++ const uint64_t ne01 = src0->ne[1]; ++ const uint64_t ne02 = src0->ne[2]; ++ const uint64_t ne03 = src0->ne[3]; ++ const uint64_t ne0 = ne00 * ne01; ++ ++ const bool use_src1 = src1 != nullptr; ++ const uint64_t ne10 = use_src1 ? src1->ne[0] : 0; ++ const uint64_t ne11 = use_src1 ? src1->ne[1] : 0; ++ const uint64_t ne12 = use_src1 ? src1->ne[2] : 0; ++ const uint64_t ne13 = use_src1 ? src1->ne[3] : 0; ++ const uint64_t ne1 = ne10 * ne11; ++ // const uint64_t nb10 = use_src1 ? src1->nb[0] : 0; ++ ++ const bool use_src2 = src2 != nullptr; ++ const uint64_t ne20 = use_src2 ? src2->ne[0] : 0; ++ const uint64_t ne21 = use_src2 ? src2->ne[1] : 0; ++ const uint64_t ne22 = use_src2 ? src2->ne[2] : 0; ++ const uint64_t ne23 = use_src2 ? src2->ne[3] : 0; ++ const uint64_t ne2 = ne20 * ne21; ++ ++ const uint64_t ned0 = dst->ne[0]; ++ const uint64_t ned1 = dst->ne[1]; ++ const uint64_t ned2 = dst->ne[2]; ++ const uint64_t ned3 = dst->ne[3]; ++ const uint64_t ned = ned0 * ned1; ++ ++ init_pushconst_fastdiv(pc); ++ ++ vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, src2, dst, op); ++ ++ if (pipeline == nullptr) { ++ std::cerr << "ggml_vulkan: Error: Missing op: " << ggml_op_name(op) << " for " << ggml_type_name(src0->type); ++ if (src1 != nullptr) { ++ std::cerr << " and " << ggml_type_name(src1->type); ++ } ++ std::cerr << " to " << ggml_type_name(dst->type) << std::endl; ++ GGML_ABORT("fatal error"); ++ } ++ ++ if (dryrun) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); ++ return; ++ } ++ ++ const bool op_supports_incontiguous = ggml_vk_op_supports_incontiguous(op); ++ ++ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ++ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; ++ ggml_backend_vk_buffer_context * src1_buf_ctx = use_src1 ? (ggml_backend_vk_buffer_context *)src1->buffer->context : nullptr; ++ ggml_backend_vk_buffer_context * src2_buf_ctx = use_src2 ? (ggml_backend_vk_buffer_context *)src2->buffer->context : nullptr; ++ ++ vk_buffer d_X = nullptr; ++ size_t x_buf_offset = 0; ++ vk_buffer d_Y = nullptr; ++ size_t y_buf_offset = 0; ++ vk_buffer d_Z = nullptr; ++ size_t z_buf_offset = 0; ++ ++ bool src0_uma = false; ++ bool src1_uma = false; ++ bool src2_uma = false; ++ ++ if (ctx->device->uma) { ++ ggml_vk_host_get(ctx->device, src0->data, d_X, x_buf_offset); ++ src0_uma = d_X != nullptr; ++ if (use_src1) { ++ ggml_vk_host_get(ctx->device, src1->data, d_Y, y_buf_offset); ++ src1_uma = d_Y != nullptr; ++ } ++ if (use_src2) { ++ ggml_vk_host_get(ctx->device, src2->data, d_Z, z_buf_offset); ++ src2_uma = d_Z != nullptr; ++ } ++ } ++ ++ uint64_t x_sz = ggml_type_size(src0->type)/ggml_blck_size(src0->type) * ne0; ++ uint64_t y_sz = use_src1 ? ggml_type_size(src1->type) * ne1 : 0; ++ uint64_t z_sz = use_src2 ? ggml_type_size(src2->type) * ne2 : 0; ++ uint64_t d_sz = ggml_type_size(dst->type) * ned; ++ ++ vk_buffer d_D = dst_buf_ctx->dev_buffer; ++ ++ // Workaround for tiny tensor inputs on ROPE ++ if (op == GGML_OP_ROPE && use_src1 && y_sz > d_D->size) { ++ y_sz = VK_WHOLE_SIZE; ++ } ++ ++ GGML_ASSERT(d_D != nullptr); ++ uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; ++ if(!src0_uma) { ++ d_X = src0_buf_ctx->dev_buffer; ++ x_buf_offset = vk_tensor_offset(src0) + src0->view_offs; ++ GGML_ASSERT(d_X != nullptr); ++ } ++ if (use_src1 && !src1_uma) { ++ d_Y = src1_buf_ctx->dev_buffer; ++ y_buf_offset = vk_tensor_offset(src1) + src1->view_offs; ++ GGML_ASSERT(d_Y != nullptr); ++ } ++ if (use_src2 && !src2_uma) { ++ d_Z = src2_buf_ctx->dev_buffer; ++ z_buf_offset = vk_tensor_offset(src2) + src2->view_offs; ++ GGML_ASSERT(d_Z != nullptr); ++ } ++ // Compute misalignment offset for descriptors and store it in in push constants, then align the descriptor offsets. ++ init_pushconst_tensor_offsets(ctx, pc, src0, src1, src2, dst); ++ x_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); ++ y_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); ++ z_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); ++ d_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); ++ ++ if (op_supports_incontiguous) { ++ x_sz = ggml_nbytes(src0); ++ y_sz = use_src1 ? ggml_nbytes(src1) : 0; ++ z_sz = use_src2 ? ggml_nbytes(src2) : 0; ++ d_sz = ggml_nbytes(dst); ++ ++ if (x_buf_offset + x_sz >= d_X->size) { ++ x_sz = VK_WHOLE_SIZE; ++ } ++ if (use_src1 && y_buf_offset + y_sz >= d_Y->size) { ++ y_sz = VK_WHOLE_SIZE; ++ } ++ if (use_src2 && z_buf_offset + z_sz >= d_Z->size) { ++ z_sz = VK_WHOLE_SIZE; ++ } ++ if (d_buf_offset + d_sz >= d_D->size) { ++ d_sz = VK_WHOLE_SIZE; ++ } ++ } ++ ++ std::array elements; ++ ++ // Single call if dimension 2 is contiguous ++ GGML_ASSERT(op_supports_incontiguous || (ggml_is_contiguous(src0) && (src1 == nullptr || ggml_is_contiguous(src1)))); ++ ++ switch (op) { ++ case GGML_OP_NORM: ++ case GGML_OP_RMS_NORM: ++ case GGML_OP_SOFT_MAX: ++ case GGML_OP_SUM_ROWS: ++ { ++ const uint32_t nr = ggml_nrows(src0); ++ if (nr > 262144) { ++ elements = { 512, 512, CEIL_DIV(nr, 262144) }; ++ } else if (nr > 512) { ++ elements = { 512, CEIL_DIV(nr, 512), 1 }; ++ } else { ++ elements = { nr, 1, 1 }; ++ } ++ } break; ++ case GGML_OP_GROUP_NORM: ++ { ++ const uint32_t num_groups = dst->op_params[0]; ++ elements = { num_groups * (uint32_t)src0->ne[3], 1, 1 }; ++ } break; ++ case GGML_OP_DIAG_MASK_INF: ++ case GGML_OP_ROPE: ++ elements = { (uint32_t)ggml_nrows(src0), (uint32_t)ne00, 1 }; ++ break; ++ case GGML_OP_GET_ROWS: ++ elements = { (uint32_t)ne00, (uint32_t)ne10, (uint32_t)(ne11 * ne12) }; ++ break; ++ case GGML_OP_ARGSORT: ++ elements = { (uint32_t)ne00, (uint32_t)ggml_nrows(src0), 1 }; ++ break; ++ case GGML_OP_IM2COL: ++ { ++ const bool is_2D = dst->op_params[6] == 1; ++ ++ const uint32_t IC = src1->ne[is_2D ? 2 : 1]; ++ ++ const uint32_t KH = is_2D ? src0->ne[1] : 1; ++ const uint32_t KW = src0->ne[0]; ++ ++ const uint32_t OH = is_2D ? dst->ne[2] : 1; ++ const uint32_t OW = dst->ne[1]; ++ ++ const uint32_t batch = src1->ne[is_2D ? 3 : 2]; ++ ++ elements = { OW * KW * KH, OH, batch * IC }; ++ } break; ++ case GGML_OP_TIMESTEP_EMBEDDING: ++ { ++ const uint32_t dim = dst->op_params[0]; ++ uint32_t half_ceil = (dim + 1) / 2; ++ elements = { half_ceil, (uint32_t)src0->ne[0], 1 }; ++ } break; ++ case GGML_OP_POOL_2D: ++ { ++ const uint32_t N = dst->ne[3]; ++ const uint32_t OC = dst->ne[2]; ++ const uint32_t OH = dst->ne[1]; ++ const uint32_t OW = dst->ne[0]; ++ elements = { N * OC * OH * OW, 1, 1}; ++ } break; ++ case GGML_OP_ADD: ++ case GGML_OP_DIV: ++ case GGML_OP_MUL: ++ case GGML_OP_SCALE: ++ case GGML_OP_SQR: ++ case GGML_OP_SIN: ++ case GGML_OP_COS: ++ case GGML_OP_CLAMP: ++ case GGML_OP_PAD: ++ case GGML_OP_REPEAT: ++ case GGML_OP_CPY: ++ case GGML_OP_CONCAT: ++ case GGML_OP_UPSCALE: ++ case GGML_OP_UNARY: ++ { ++ const uint32_t ne = ggml_nelements(dst); ++ if (ne > 262144) { ++ elements = { 512, 512, CEIL_DIV(ne, 262144) }; ++ } else if (ne > 512) { ++ elements = { 512, CEIL_DIV(ne, 512), 1 }; ++ } else { ++ elements = { ne, 1, 1 }; ++ } ++ } break; ++ default: ++ elements = { (uint32_t)ggml_nelements(src0), 1, 1 }; ++ break; ++ } ++ ++ if (!op_supports_incontiguous) { ++ if (x_sz != VK_WHOLE_SIZE) { ++ x_sz *= ne02 * ne03; ++ } ++ if (use_src1 && y_sz != VK_WHOLE_SIZE) { ++ y_sz *= ne12 * ne13; ++ } ++ if (use_src2 && z_sz != VK_WHOLE_SIZE) { ++ z_sz *= ne22 * ne23; ++ } ++ if (d_sz != VK_WHOLE_SIZE) { ++ d_sz *= ned2 * ned3; ++ } ++ } ++ ++ if (op == GGML_OP_SOFT_MAX) { ++ // Empty src1 is possible in soft_max, but the shader needs a buffer ++ vk_subbuffer subbuf_y; ++ if (use_src1) { ++ subbuf_y = { d_Y, y_buf_offset, y_sz }; ++ } else { ++ subbuf_y = { d_X, 0, x_sz }; ++ } ++ ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, subbuf_y, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); ++ } else if (op == GGML_OP_ROPE) { ++ // Empty src2 is possible in rope, but the shader needs a buffer ++ vk_subbuffer subbuf_z; ++ if (use_src2) { ++ subbuf_z = { d_Z, z_buf_offset, z_sz }; ++ } else { ++ subbuf_z = { d_X, 0, x_sz }; ++ } ++ ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, subbuf_z, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); ++ } else if (op == GGML_OP_IM2COL) { ++ // im2col uses only src1 and dst buffers ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); ++ } else if (use_src2) { ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_Z, z_buf_offset, z_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); ++ } else if (use_src1) { ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); ++ } else { ++ ggml_vk_sync_buffers(subctx); ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); ++ } ++} ++ ++static void ggml_vk_get_rows(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t src1_type_size = ggml_type_size(src1->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_GET_ROWS, { ++ (uint32_t)ggml_nelements(src0), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_acc(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t src1_type_size = ggml_type_size(src1->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ int nb1 = dst->op_params[0] / 4; // 4 bytes of float32 ++ int nb2 = dst->op_params[1] / 4; // 4 bytes of float32 ++ // int nb3 = dst->op_params[2] / 4; // 4 bytes of float32 - unused ++ int offset = dst->op_params[3] / 4; // offset in bytes ++ ++ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_ACC, { ++ (uint32_t)ggml_nelements(src0), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)nb1, (uint32_t)nb2, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t)nb1, (uint32_t)nb2, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, offset, ++ }, dryrun); ++} ++ ++static void ggml_vk_add(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t src1_type_size = ggml_type_size(src1->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_ADD, { ++ (uint32_t)ggml_nelements(src0), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_mul(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t src1_type_size = ggml_type_size(src1->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_MUL, { ++ (uint32_t)ggml_nelements(src0), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_div(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t src1_type_size = ggml_type_size(src1->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_DIV, { ++ (uint32_t)ggml_nelements(src0), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_op_f32_rwkv6(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, const vk_op_rwkv_wkv6_push_constants&& pc, bool dryrun = false) { ++ const ggml_tensor * k = dst->src[0]; ++ const ggml_tensor * v = dst->src[1]; ++ const ggml_tensor * r = dst->src[2]; ++ const ggml_tensor * tf = dst->src[3]; ++ const ggml_tensor * td = dst->src[4]; ++ const ggml_tensor * state = dst->src[5]; ++ ++ GGML_ASSERT(!ggml_is_quantized(k->type)); ++ GGML_ASSERT(!ggml_is_quantized(v->type)); ++ GGML_ASSERT(!ggml_is_quantized(r->type)); ++ GGML_ASSERT(!ggml_is_quantized(tf->type)); ++ GGML_ASSERT(!ggml_is_quantized(td->type)); ++ GGML_ASSERT(!ggml_is_quantized(state->type)); ++ GGML_ASSERT(dst->buffer != nullptr); ++ ++ vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, k, v, r, dst, GGML_OP_RWKV_WKV6); ++ GGML_ASSERT(pipeline != nullptr); ++ ++ if (dryrun) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); ++ return; ++ } ++ ++ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ++ ggml_backend_vk_buffer_context * k_buf_ctx = (ggml_backend_vk_buffer_context *)k->buffer->context; ++ ggml_backend_vk_buffer_context * v_buf_ctx = (ggml_backend_vk_buffer_context *)v->buffer->context; ++ ggml_backend_vk_buffer_context * r_buf_ctx = (ggml_backend_vk_buffer_context *)r->buffer->context; ++ ggml_backend_vk_buffer_context * tf_buf_ctx = (ggml_backend_vk_buffer_context *)tf->buffer->context; ++ ggml_backend_vk_buffer_context * td_buf_ctx = (ggml_backend_vk_buffer_context *)td->buffer->context; ++ ggml_backend_vk_buffer_context * state_buf_ctx = (ggml_backend_vk_buffer_context *)state->buffer->context; ++ ++ ggml_vk_sync_buffers(subctx); ++ ++ vk_buffer d_D = nullptr, d_K = nullptr, d_V = nullptr, d_R = nullptr, d_TF = nullptr, d_TD = nullptr, d_State = nullptr; ++ size_t k_offset = 0, v_offset = 0, r_offset = 0, tf_offset = 0, td_offset = 0, state_offset = 0, dst_offset = 0; ++ bool K_uma = false, V_uma = false, R_uma = false, TF_uma = false, TD_uma = false, STATE_uma = false, DST_uma = false; ++ ++ if (ctx->device->uma) { ++ ggml_vk_host_get(ctx->device, k->data, d_K, k_offset); ++ ggml_vk_host_get(ctx->device, v->data, d_V, v_offset); ++ ggml_vk_host_get(ctx->device, r->data, d_R, r_offset); ++ ggml_vk_host_get(ctx->device, tf->data, d_TF, tf_offset); ++ ggml_vk_host_get(ctx->device, td->data, d_TD, td_offset); ++ ggml_vk_host_get(ctx->device, state->data, d_State, state_offset); ++ ggml_vk_host_get(ctx->device, dst->data, d_D, dst_offset); ++ ++ K_uma = d_K != nullptr; ++ V_uma = d_V != nullptr; ++ R_uma = d_R != nullptr; ++ TF_uma = d_TF != nullptr; ++ TD_uma = d_TD != nullptr; ++ STATE_uma = d_State != nullptr; ++ DST_uma = d_D != nullptr; ++ } ++ ++ if (!K_uma) { ++ d_K = k_buf_ctx->dev_buffer; ++ k_offset = vk_tensor_offset(k) + k->view_offs; ++ } ++ if (!V_uma) { ++ d_V = v_buf_ctx->dev_buffer; ++ v_offset = vk_tensor_offset(v) + v->view_offs; ++ } ++ if (!R_uma) { ++ d_R = r_buf_ctx->dev_buffer; ++ r_offset = vk_tensor_offset(r) + r->view_offs; ++ } ++ if (!TF_uma) { ++ d_TF = tf_buf_ctx->dev_buffer; ++ tf_offset = vk_tensor_offset(tf) + tf->view_offs; ++ } ++ if (!TD_uma) { ++ d_TD = td_buf_ctx->dev_buffer; ++ td_offset = vk_tensor_offset(td) + td->view_offs; ++ } ++ if (!STATE_uma) { ++ d_State = state_buf_ctx->dev_buffer; ++ state_offset = vk_tensor_offset(state) + state->view_offs; ++ } ++ if (!DST_uma) { ++ d_D = dst_buf_ctx->dev_buffer; ++ dst_offset = vk_tensor_offset(dst) + dst->view_offs; ++ } ++ ++ const uint64_t k_size = ggml_nbytes(k); ++ const uint64_t v_size = ggml_nbytes(v); ++ const uint64_t r_size = ggml_nbytes(r); ++ const uint64_t tf_size = ggml_nbytes(tf); ++ const uint64_t td_size = ggml_nbytes(td); ++ const uint64_t state_size = ggml_nbytes(state); ++ const uint64_t dst_size = ggml_nbytes(dst); ++ ++ std::array elements = { ++ (uint32_t)(pc.B * pc.H), ++ 1, ++ 1 ++ }; ++ ++ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { ++ vk_subbuffer{ d_K, k_offset, k_size }, ++ vk_subbuffer{ d_V, v_offset, v_size }, ++ vk_subbuffer{ d_R, r_offset, r_size }, ++ vk_subbuffer{ d_TF, tf_offset, tf_size }, ++ vk_subbuffer{ d_TD, td_offset, td_size }, ++ vk_subbuffer{ d_State, state_offset, state_size }, ++ vk_subbuffer{ d_D, dst_offset, dst_size } ++ }, sizeof(vk_op_rwkv_wkv6_push_constants), &pc, elements); ++} ++ ++static void ggml_vk_rwkv_wkv6(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, bool dryrun = false) { ++ const size_t seq_length = dst->src[0]->ne[3]; ++ const size_t n_embed = dst->ne[0]; ++ const size_t n_heads = dst->src[0]->ne[2]; ++ const size_t n_seqs = dst->src[5]->ne[1]; ++ ++ ggml_vk_op_f32_rwkv6( ++ ctx, subctx, dst, ++ { ++ (uint32_t)n_seqs, ++ (uint32_t)seq_length, ++ (uint32_t)n_embed, ++ (uint32_t)n_heads, ++ }, ++ dryrun ++ ); ++} ++ ++static void ggml_vk_concat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ int * op_params = (int *)dst->op_params; ++ ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t src1_type_size = ggml_type_size(src1->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_CONCAT, { ++ (uint32_t)ggml_nelements(dst), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, op_params[0], ++ }, dryrun); ++} ++ ++static void ggml_vk_upscale(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ ++ const float sf0 = (float)dst->ne[0] / src0->ne[0]; ++ const float sf1 = (float)dst->ne[1] / src0->ne[1]; ++ const float sf2 = (float)dst->ne[2] / src0->ne[2]; ++ const float sf3 = (float)dst->ne[3] / src0->ne[3]; ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_UPSCALE, { ++ (uint32_t)ggml_nelements(dst), 0, 0, ++ (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t)dst->ne[0], (uint32_t)dst->ne[1], (uint32_t)dst->ne[2],(uint32_t)dst->ne[3], ++ sf0, sf1, sf2, sf3, ++ }, dryrun); ++} ++ ++static void ggml_vk_scale(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ float * op_params = (float *)dst->op_params; ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SCALE, { ++ (uint32_t)ggml_nelements(src0), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ op_params[0], 0.0f, ++ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_sqr(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SQR, { ++ (uint32_t)ggml_nelements(src0), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, ++ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_sin(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SIN, { ++ (uint32_t)ggml_nelements(src0), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, ++ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_cos(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_COS, { ++ (uint32_t)ggml_nelements(src0), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, ++ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_clamp(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ float * op_params = (float *)dst->op_params; ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_CLAMP, { ++ (uint32_t)ggml_nelements(src0), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ op_params[0], op_params[1], ++ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_pad(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_PAD, { ++ (uint32_t)ggml_nelements(dst), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, ++ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_repeat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_REPEAT, { ++ (uint32_t)ggml_nelements(dst), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, ++ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_cpy(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t src0_type_size = ggml_type_size(src0->type); ++ const uint32_t dst_type_size = ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_CPY, { ++ (uint32_t)ggml_nelements(src0), ++ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, ++ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, ++ 0, ++ 0.0f, 0.0f, ++ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ++ }, dryrun); ++} ++ ++static void ggml_vk_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ float * op_params = (float *)dst->op_params; ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_NORM, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0], 0.0f }, dryrun); ++} ++ ++static void ggml_vk_group_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ const int * int_op_params = (const int *)dst->op_params; ++ const float * float_op_params = (const float *)dst->op_params; ++ ++ const uint32_t num_groups = int_op_params[0]; ++ const float eps = float_op_params[1]; ++ const uint32_t group_size = src0->ne[0] * src0->ne[1] * ((src0->ne[2] + num_groups - 1) / num_groups); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_GROUP_NORM, { group_size, 0, eps, 0.0f }, dryrun); ++} ++ ++static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ float * op_params = (float *)dst->op_params; ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_RMS_NORM, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0], 0.0f }, dryrun); ++} ++ ++static void ggml_vk_unary(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_UNARY, { (uint32_t)ggml_nelements(src0), 0, 0.0f, 0.0f }, dryrun); ++} ++ ++static void ggml_vk_diag_mask_inf(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ int32_t * op_params = (int32_t *)dst->op_params; ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_DIAG_MASK_INF, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0] }, dryrun); ++} ++ ++static void ggml_vk_soft_max(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ float * op_params = (float *)dst->op_params; ++ ++ float scale = op_params[0]; ++ float max_bias = op_params[1]; ++ ++ const uint32_t ncols = (uint32_t)src0->ne[0]; ++ const uint32_t nrows_x = (uint32_t)ggml_nrows(src0); ++ const uint32_t nrows_y = (uint32_t)src0->ne[1]; ++ ++ const uint32_t n_head_kv = nrows_x/nrows_y; ++ const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head_kv)); ++ ++ const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); ++ const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_SOFT_MAX, { ++ ncols, ++ src1 != nullptr ? nrows_y : (uint32_t)0, ++ scale, max_bias, ++ m0, m1, ++ n_head_log2, ++ nrows_x, ++ }, dryrun); ++} ++ ++static void ggml_vk_rope(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, bool dryrun = false) { ++ const int n_dims = ((int32_t *) dst->op_params)[1]; ++ // const int mode = ((int32_t *) dst->op_params)[2]; ++ // const int n_ctx = ((int32_t *) dst->op_params)[3]; ++ const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; ++ const float freq_base = ((float *) dst->op_params)[5]; ++ const float freq_scale = ((float *) dst->op_params)[6]; ++ const float ext_factor = ((float *) dst->op_params)[7]; ++ const float attn_factor = ((float *) dst->op_params)[8]; ++ const float beta_fast = ((float *) dst->op_params)[9]; ++ const float beta_slow = ((float *) dst->op_params)[10]; ++ ++ float corr_dims[2]; ++ ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); ++ ++ const float theta_scale = powf(freq_base, -2.0f/n_dims); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, src1, src2, dst, GGML_OP_ROPE, { ++ (uint32_t)src0->ne[0], (uint32_t)n_dims, freq_scale, (uint32_t)src0->ne[1], ++ freq_base, ext_factor, attn_factor, {corr_dims[0], corr_dims[1]}, theta_scale, ++ src2 != nullptr, ++ }, dryrun); ++} ++ ++static void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ int32_t * op_params = (int32_t *)dst->op_params; ++ ++ uint32_t ncols = src0->ne[0]; ++ ++ uint32_t ncols_pad = 1; ++ while (ncols_pad < ncols) { ++ ncols_pad *= 2; ++ } ++ ++ GGML_ASSERT(ncols_pad <= 1024); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_ARGSORT, { ++ ncols, ++ ncols_pad, ++ op_params[0], ++ }, dryrun); ++} ++ ++static void ggml_vk_sum_rows(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SUM_ROWS, { (uint32_t)src0->ne[0], 0, 0.0f, 0.0f }, dryrun); ++} ++ ++static void ggml_vk_im2col(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { ++ const int32_t s0 = dst->op_params[0]; ++ const int32_t s1 = dst->op_params[1]; ++ const int32_t p0 = dst->op_params[2]; ++ const int32_t p1 = dst->op_params[3]; ++ const int32_t d0 = dst->op_params[4]; ++ const int32_t d1 = dst->op_params[5]; ++ ++ const bool is_2D = dst->op_params[6] == 1; ++ ++ const uint32_t IC = src1->ne[is_2D ? 2 : 1]; ++ const uint32_t IH = is_2D ? src1->ne[1] : 1; ++ const uint32_t IW = src1->ne[0]; ++ ++ const uint32_t KH = is_2D ? src0->ne[1] : 1; ++ const uint32_t KW = src0->ne[0]; ++ ++ const uint32_t OH = is_2D ? dst->ne[2] : 1; ++ const uint32_t OW = dst->ne[1]; ++ ++ const uint32_t offset_delta = src1->nb[is_2D ? 2 : 1] / 4; // nb is byte offset, src is type float32 ++ const uint32_t batch_offset = src1->nb[is_2D ? 3 : 2] / 4; // nb is byte offset, src is type float32 ++ ++ const uint32_t pelements = OW * KW * KH; ++ ++ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_IM2COL, { ++ batch_offset, offset_delta, ++ IC, IW, IH, OW, OH, KW, KH, ++ pelements, ++ IC * KH * KW, ++ s0, s1, p0, p1, d0, d1, ++ }, dryrun); ++} ++ ++static void ggml_vk_timestep_embedding(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ const uint32_t dim = dst->op_params[0]; ++ const uint32_t max_period = dst->op_params[1]; ++ const uint32_t nb1 = dst->nb[1] / ggml_type_size(dst->type); ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_TIMESTEP_EMBEDDING, { ++ nb1, dim, max_period, ++ }, dryrun); ++} ++ ++static void ggml_vk_pool_2d(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ uint32_t op = static_cast(dst->op_params[0]); ++ const int32_t k1 = dst->op_params[1]; ++ const int32_t k0 = dst->op_params[2]; ++ const int32_t s1 = dst->op_params[3]; ++ const int32_t s0 = dst->op_params[4]; ++ const int32_t p1 = dst->op_params[5]; ++ const int32_t p0 = dst->op_params[6]; ++ ++ const uint32_t IH = src0->ne[1]; ++ const uint32_t IW = src0->ne[0]; ++ ++ const uint32_t N = dst->ne[3]; ++ ++ const uint32_t OC = dst->ne[2]; ++ const uint32_t OH = dst->ne[1]; ++ const uint32_t OW = dst->ne[0]; ++ ++ const uint32_t parallel_elements = N * OC * OH * OW; ++ ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_POOL_2D, { ++ IW, IH, OW, OH, OC, ++ parallel_elements, ++ op, ++ k0, k1, s0, s1, p0, p1, ++ }, dryrun); ++} ++ ++static void ggml_vk_leaky_relu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { ++ const float * op_params = (const float *)dst->op_params; ++ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_LEAKY_RELU, { (uint32_t)ggml_nelements(src0), 0, op_params[0], 0.0f }, dryrun); ++} ++ ++#ifdef GGML_VULKAN_RUN_TESTS ++static void ggml_vk_print_matrix_area(const void * data, ggml_type type, int ne0, int ne1, int i0, int i1, int i2) { ++ if (type != GGML_TYPE_F32 && type != GGML_TYPE_F16) { ++ return; ++ } ++ i0 = std::max(i0, 5); ++ i1 = std::max(i1, 5); ++ i2 = std::max(i2, 0); ++ fprintf(stderr, " "); ++ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { ++ fprintf(stderr, "%7d ", idx1); ++ } ++ fprintf(stderr, "\n"); ++ for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { ++ fprintf(stderr, "%7d: ", idx0); ++ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { ++ if (idx0 >= 0 && idx0 < ne0 && idx1 >= 0 && idx1 < ne1) { ++ float val; ++ if (type == GGML_TYPE_F32) { ++ val = *((const float *) data + i2*ne1*ne0 + idx1*ne0 + idx0); ++ } else if (type == GGML_TYPE_F16) { ++ val = ggml_fp16_to_fp32(*((const ggml_fp16_t *) data + i2*ne1*ne0 + idx1*ne0 + idx0)); ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ fprintf(stderr, "% 7.2f ", val); ++ } else { ++ fprintf(stderr, " "); ++ } ++ } ++ fprintf(stderr, "\n"); ++ } ++} ++ ++template ++static void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, int split_k, int shader_size) { ++ VK_LOG_DEBUG("ggml_vk_test_matmul(" << m << ", " << n << ", " << k << ", " << batch << ", " << num_it << ", " << split_k << ", " << shader_size << ")"); ++ const size_t x_ne = m * k * batch; ++ const size_t y_ne = k * n * batch; ++ const size_t d_ne = m * n * batch; ++ ++ vk_pipeline p; ++ std::string shname; ++ if (shader_size == 0) { ++ if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32->a_s; ++ shname = "F32_ALIGNED_S"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32_f16->a_s; ++ shname = "F32_F16_ALIGNED_S"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_s; ++ shname = "F16_F32_ALIGNED_S"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16.f32acc->a_s; ++ shname = "F16_ALIGNED_S"; ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ } else if (shader_size == 1) { ++ if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32->a_m; ++ shname = "F32_ALIGNED_M"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32_f16->a_m; ++ shname = "F32_F16_ALIGNED_M"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_m; ++ shname = "F16_F32_ALIGNED_M"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16.f32acc->a_m; ++ shname = "F16_ALIGNED_M"; ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ } else if (shader_size == 2) { ++ if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32->a_l; ++ shname = "F32_ALIGNED_L"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32_f16->a_l; ++ shname = "F32_F16_ALIGNED_L"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_l; ++ shname = "F16_F32_ALIGNED_L"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16.f32acc->a_l; ++ shname = "F16_ALIGNED_L"; ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ } else { ++ GGML_ASSERT(0); ++ } ++ ++ const size_t kpad = ggml_vk_align_size(k, p->align); ++ ++ if (k != kpad) { ++ if (shader_size == 0) { ++ if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32->s; ++ shname = "F32_S"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32_f16->s; ++ shname = "F32_F16_S"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16_f32.f32acc->s; ++ shname = "F16_F32_S"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16.f32acc->s; ++ shname = "F16_S"; ++ } ++ } else if (shader_size == 1) { ++ if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32->m; ++ shname = "F32_M"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32_f16->m; ++ shname = "F32_F16_M"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16_f32.f32acc->m; ++ shname = "F16_F32_M"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16.f32acc->m; ++ shname = "F16_M"; ++ } ++ } else if (shader_size == 2) { ++ if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32->l; ++ shname = "F32_L"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f32_f16->l; ++ shname = "F32_F16_L"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16_f32.f32acc->l; ++ shname = "F16_F32_L"; ++ } else if (std::is_same() && std::is_same()) { ++ p = ctx->device->pipeline_matmul_f16.f32acc->l; ++ shname = "F16_L"; ++ } ++ } ++ } ++ ++ ggml_pipeline_request_descriptor_sets(ctx->device, p, num_it); ++ if (split_k > 1) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, num_it); ++ ++ if (ctx->prealloc_split_k == nullptr || ctx->prealloc_split_k->size < sizeof(float) * d_ne * split_k) { ++ // Resize buffer ++ if (ctx->prealloc_split_k != nullptr) { ++ ggml_vk_destroy_buffer(ctx->prealloc_split_k); ++ } ++ ctx->prealloc_split_k = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne * split_k, vk::MemoryPropertyFlagBits::eDeviceLocal); ++ } ++ } ++ ++ ggml_pipeline_allocate_descriptor_sets(ctx->device); ++ ++ vk_buffer d_X = ggml_vk_create_buffer_check(ctx->device, sizeof(X_TYPE) * x_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); ++ vk_buffer d_Y = ggml_vk_create_buffer_check(ctx->device, sizeof(Y_TYPE) * y_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); ++ vk_buffer d_D = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); ++ ++ X_TYPE* x = (X_TYPE *) malloc(sizeof(X_TYPE) * x_ne); ++ Y_TYPE* y = (Y_TYPE *) malloc(sizeof(Y_TYPE) * y_ne); ++ float* d = (float *) malloc(sizeof(float) * d_ne); ++ ++ for (size_t i = 0; i < x_ne; i++) { ++ if (std::is_same()) { ++ x[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; ++ // x[i] = 1.0f; ++ // x[i] = i + 1; ++ // x[i] = (i % k == i / k) ? 1.0f : 0.0f; ++ } else if (std::is_same()) { ++ x[i] = ggml_fp32_to_fp16((rand() / (float)RAND_MAX) * 2.0f - 1.0f); ++ // x[i] = ggml_fp32_to_fp16(1.0f); ++ // x[i] = ggml_fp32_to_fp16(i + 1); ++ // x[i] = ggml_fp32_to_fp16((i % k == i / k) ? 1.0f : 0.0f); ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ } ++ for (size_t i = 0; i < y_ne; i++) { ++ if (std::is_same()) { ++ y[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; ++ // y[i] = (i % k == i / k) ? 1.0f : 0.0f; ++ // y[i] = i + 1; ++ } else if (std::is_same()) { ++ y[i] = ggml_fp32_to_fp16((rand() / (float)RAND_MAX) * 2.0f - 1.0f); ++ // y[i] = ggml_fp32_to_fp16((i % k == i / k) ? 1.0f : 0.0f); ++ // y[i] = ggml_fp32_to_fp16(i + 1); ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ } ++ ++ ggml_vk_buffer_write(d_X, 0, x, sizeof(X_TYPE) * k * m * batch); ++ ggml_vk_buffer_write(d_Y, 0, y, sizeof(Y_TYPE) * k * n * batch); ++ ++ vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); ++ ggml_vk_ctx_begin(ctx->device, subctx); ++ for (size_t i = 0; i < num_it; i++) { ++ ggml_vk_matmul( ++ ctx, subctx, p, ggml_vk_subbuffer(d_X), ggml_vk_subbuffer(d_Y), ggml_vk_subbuffer(d_D), ggml_vk_subbuffer(ctx->prealloc_split_k), ++ m, n, k, ++ k, k, m, k*m, k*n, m*n, ++ split_k, batch, batch, batch, 1, 1 ++ ); ++ } ++ ggml_vk_ctx_end(subctx); ++ ++ auto begin = std::chrono::high_resolution_clock::now(); ++ ggml_vk_submit(subctx, ctx->fence); ++ VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_matmul waitForFences"); ++ ctx->device->device.resetFences({ ctx->fence }); ++ ++ auto end = std::chrono::high_resolution_clock::now(); ++ double time = std::chrono::duration_cast(end-begin).count() / 1000.0; ++ ++ // copy dst to host ++ ggml_vk_buffer_read(d_D, 0, d, sizeof(float) * d_ne); ++ ++ float * d_chk = (float *) malloc(sizeof(float) * d_ne); ++ ++ ggml_init_params iparams = { ++ /*.mem_size =*/ 1024*1024*1024, ++ /*.mem_buffer =*/ NULL, ++ /*.no_alloc =*/ true, ++ }; ++ ++ ggml_context * ggml_ctx = ggml_init(iparams); ++ ++ ggml_type src0_type; ++ ggml_type src1_type; ++ ++ if (std::is_same()) { ++ src0_type = GGML_TYPE_F32; ++ } else if (std::is_same()) { ++ src0_type = GGML_TYPE_F16; ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ if (std::is_same()) { ++ src1_type = GGML_TYPE_F32; ++ } else if (std::is_same()) { ++ src1_type = GGML_TYPE_F16; ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ ++ ggml_tensor * src0_ggml = ggml_new_tensor_3d(ggml_ctx, src0_type, k, m, batch); ++ ggml_tensor * src1_ggml = ggml_new_tensor_3d(ggml_ctx, src1_type, k, n, batch); ++ ggml_tensor * tensor_ggml = ggml_mul_mat(ggml_ctx, src0_ggml, src1_ggml); ++ ++ src0_ggml->data = x; ++ src1_ggml->data = y; ++ tensor_ggml->data = d_chk; ++ ++ ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); ++ ggml_build_forward_expand(cgraph, tensor_ggml); ++ ++ ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 1); ++ ++ ggml_free(ggml_ctx); ++ ++ double avg_err = 0.0; ++ int first_err_n = -1; ++ int first_err_m = -1; ++ int first_err_b = -1; ++ ++ for (size_t i = 0; i < m*n*batch; i++) { ++ double err = std::fabs(d[i] - d_chk[i]); ++ avg_err += err; ++ ++ if ((err > 0.05f || std::isnan(err)) && first_err_n == -1) { ++ first_err_b = i / (m * n); ++ first_err_n = (i % (m * n)) / m; ++ first_err_m = (i % (m * n)) % m; ++ } ++ } ++ ++ avg_err /= m * n; ++ ++ double tflops = 2.0*m*n*k*batch*num_it / (time / 1000.0) / (1000.0*1000.0*1000.0*1000.0); ++ ++ std::cerr << "TEST " << shname << " m=" << m << " n=" << n << " k=" << k << " batch=" << batch << " split_k=" << split_k << " matmul " << time / num_it << "ms " << tflops << " TFLOPS avg_err=" << avg_err << std::endl; ++ ++ if (avg_err > 0.1 || std::isnan(avg_err)) { ++ std::cerr << "m = " << first_err_m << " n = " << first_err_n << " b = " << first_err_b << std::endl; ++ std::cerr << "Actual result: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(d, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ std::cerr << "Expected result: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(d_chk, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ ++ if (split_k > 1) { ++ float * split_k_buf = (float *) malloc(sizeof(float) * d_ne * split_k); ++ ggml_vk_buffer_read(ctx->prealloc_split_k, 0, split_k_buf, sizeof(float) * d_ne * split_k); ++ ++ std::cerr << "d_buf0: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(split_k_buf, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ ++ std::cerr << "d_buf1: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(split_k_buf + d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ ++ std::cerr << "d_buf2: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(split_k_buf + 2 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ ++ std::cerr << "d_buf3: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(split_k_buf + 3 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ ++ free(split_k_buf); ++ } ++ } ++ ++ free(d_chk); ++ ++ ggml_vk_queue_cleanup(ctx->device, ctx->device->transfer_queue); ++ ggml_vk_queue_cleanup(ctx->device, ctx->device->compute_queue); ++ ++ ggml_vk_destroy_buffer(d_X); ++ ggml_vk_destroy_buffer(d_Y); ++ ggml_vk_destroy_buffer(d_D); ++ ++ ggml_pipeline_cleanup(p); ++ ggml_pipeline_cleanup(ctx->device->pipeline_matmul_split_k_reduce); ++ ++ free(x); ++ free(y); ++ free(d); ++} ++ ++static void ggml_vk_print_tensor_area(const ggml_tensor * tensor, int i0, int i1, int i2, int i3) { ++ if (tensor->type != GGML_TYPE_F32 && tensor->type != GGML_TYPE_F16) { ++ return; ++ } ++ i0 = std::max(i0, 5); ++ i1 = std::max(i1, 5); ++ i2 = std::max(i2, 0); ++ i3 = std::max(i3, 0); ++ fprintf(stderr, " "); ++ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { ++ fprintf(stderr, "%7d ", idx1); ++ } ++ fprintf(stderr, "\n"); ++ for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { ++ fprintf(stderr, "%7d: ", idx0); ++ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { ++ if (idx0 >= 0 && idx0 < tensor->ne[0] && idx1 >= 0 && idx1 < tensor->ne[1] && i2 >= 0 && i2 < tensor->ne[2] && i3 >= 0 && i3 < tensor->ne[3]) { ++ float val; ++ if (tensor->type == GGML_TYPE_F32) { ++ val = *(float *) ((char *) tensor->data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); ++ } else if (tensor->type == GGML_TYPE_F16) { ++ val = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) tensor->data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0])); ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ fprintf(stderr, "% 7.2f ", val); ++ } else { ++ fprintf(stderr, " "); ++ } ++ } ++ fprintf(stderr, "\n"); ++ } ++} ++ ++static void ggml_vk_quantize_data(const float * from, void * to, size_t ne, ggml_type quant) { ++ ggml_quantize_chunk(quant, from, to, 0, 1, ne, nullptr); ++} ++ ++static void ggml_vk_dequantize_data(const void * from, float * to, size_t ne, ggml_type quant) { ++ if (quant == GGML_TYPE_F32) { ++ memcpy(to, from, sizeof(float) * ne); ++ return; ++ } ++ ++ const auto * tt = ggml_get_type_traits(quant); ++ ++ ggml_to_float_t dequant_fn = tt->to_float; ++ ++ dequant_fn(from, to, ne); ++} ++ ++static void ggml_vk_test_dequant(ggml_backend_vk_context * ctx, size_t ne, ggml_type quant) { ++ VK_LOG_DEBUG("ggml_vk_test_dequant(" << ne << ")"); ++ const size_t x_sz = sizeof(float) * ne; ++ const size_t x_sz_f16 = sizeof(ggml_fp16_t) * ne; ++ const size_t qx_sz = ne * ggml_type_size(quant)/ggml_blck_size(quant); ++ float * x = (float *) malloc(x_sz); ++ void * qx = malloc(qx_sz); ++ vk_buffer qx_buf = ggml_vk_create_buffer_check(ctx->device, qx_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); ++ vk_buffer x_buf = ggml_vk_create_buffer_check(ctx->device, x_sz_f16, vk::MemoryPropertyFlagBits::eDeviceLocal); ++ float * x_ref = (float *) malloc(x_sz); ++ ggml_fp16_t * x_chk = (ggml_fp16_t *) malloc(x_sz_f16); ++ ++ for (size_t i = 0; i < ne; i++) { ++ x[i] = rand() / (float)RAND_MAX; ++ } ++ ++ vk_pipeline p = ggml_vk_get_to_fp16(ctx, quant); ++ ++ ggml_vk_quantize_data(x, qx, ne, quant); ++ ggml_vk_dequantize_data(qx, x_ref, ne, quant); ++ ++ ggml_pipeline_request_descriptor_sets(ctx->device, p, 1); ++ ++ ggml_pipeline_allocate_descriptor_sets(ctx->device); ++ ++ ggml_vk_buffer_write(qx_buf, 0, qx, qx_sz); ++ ++ vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); ++ ggml_vk_ctx_begin(ctx->device, subctx); ++ const std::vector pc = { 1, (uint32_t)ne, (uint32_t)ne, (uint32_t)ne, (uint32_t)ne }; ++ ggml_vk_dispatch_pipeline(ctx, subctx, p, { vk_subbuffer{ qx_buf, 0, qx_sz }, vk_subbuffer{ x_buf, 0, x_sz_f16 } }, pc.size() * sizeof(int), pc.data(), { (uint32_t)ne, 1, 1}); ++ ggml_vk_ctx_end(subctx); ++ ++ auto begin = std::chrono::high_resolution_clock::now(); ++ ++ ggml_vk_submit(subctx, ctx->fence); ++ VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences"); ++ ctx->device->device.resetFences({ ctx->fence }); ++ ++ auto end = std::chrono::high_resolution_clock::now(); ++ ++ double ms_dequant = std::chrono::duration_cast(end-begin).count() / 1000.0; ++ ggml_vk_buffer_read(x_buf, 0, x_chk, x_sz_f16); ++ ++ int first_err = -1; ++ ++ double avg_err = 0.0; ++ for (size_t i = 0; i < ne; i++) { ++ double error = std::fabs(x_ref[i] - ggml_fp16_to_fp32(x_chk[i])); ++ avg_err += error; ++ ++ if (first_err < 0 && error > 0.05) { ++ first_err = i; ++ } ++ } ++ ++ avg_err /= ne; ++ ++ std::cerr << "TEST DEQUANT " << ggml_type_name(quant) << " time=" << ms_dequant << "ms avg_err=" << avg_err << std::endl; ++ ++ if (avg_err > 0.1) { ++ std::cerr << "first_error = " << first_err << std::endl; ++ std::cerr << "Actual result: " << std::endl << std::endl; ++ for (int i = std::max(0, first_err - 5); i < std::min((int)ne, first_err + 5); i++) { ++ std::cerr << ggml_fp16_to_fp32(x_chk[i]) << ", "; ++ } ++ std::cerr << std::endl << "Expected result: " << std::endl << std::endl; ++ for (int i = std::max(0, first_err - 5); i < std::min((int)ne, first_err + 5); i++) { ++ std::cerr << x_ref[i] << ", "; ++ } ++ std::cerr << std::endl; ++ } ++ ++ ggml_vk_destroy_buffer(x_buf); ++ ggml_vk_destroy_buffer(qx_buf); ++ ++ free(x); ++ free(qx); ++ free(x_ref); ++ free(x_chk); ++} ++ ++static void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, size_t split_k, size_t shader_size, ggml_type quant) { ++ VK_LOG_DEBUG("ggml_vk_test_dequant_matmul(" << m << ", " << n << ", " << k << ", " << batch << ", " << num_it << ", " << split_k << ", " << ggml_type_name(quant) << ")"); ++ const size_t x_ne = m * k * batch; ++ const size_t y_ne = k * n * batch; ++ const size_t d_ne = m * n * batch; ++ ++ vk_pipeline p; ++ std::string shname; ++ if (shader_size == 0) { ++ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_s : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_s; ++ shname = std::string(ggml_type_name(quant)) + "_ALIGNED_S"; ++ } else if (shader_size == 1) { ++ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_m : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_m; ++ shname = std::string(ggml_type_name(quant)) + "_ALIGNED_M"; ++ } else if (shader_size == 2) { ++ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_l : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_l; ++ shname = std::string(ggml_type_name(quant)) + "_ALIGNED_L"; ++ } else { ++ GGML_ASSERT(0); ++ } ++ ++ const size_t kpad = ggml_vk_align_size(k, p->align); ++ ++ if (k != kpad) { ++ if (shader_size == 0) { ++ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->s : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->s; ++ shname = std::string(ggml_type_name(quant)) + "_S"; ++ } else if (shader_size == 1) { ++ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->m : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->m; ++ shname = std::string(ggml_type_name(quant)) + "_M"; ++ } else if (shader_size == 2) { ++ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->l : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->l; ++ shname = std::string(ggml_type_name(quant)) + "_L"; ++ } else { ++ GGML_ASSERT(0); ++ } ++ } ++ ++ const size_t x_sz = sizeof(float) * x_ne; ++ const size_t y_sz = sizeof(float) * y_ne; ++ const size_t qx_sz = x_ne * ggml_type_size(quant)/ggml_blck_size(quant); ++ const size_t d_sz = sizeof(float) * d_ne; ++ float * x = (float *) malloc(x_sz); ++ float * y = (float *) malloc(y_sz); ++ void * qx = malloc(qx_sz); ++ vk_buffer qx_buf = ggml_vk_create_buffer_check(ctx->device, qx_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); ++ vk_buffer y_buf = ggml_vk_create_buffer_check(ctx->device, y_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); ++ vk_buffer d_buf = ggml_vk_create_buffer_check(ctx->device, d_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); ++ float * d = (float *) malloc(d_sz); ++ float * d_chk = (float *) malloc(d_sz); ++ ++ for (size_t i = 0; i < x_ne; i++) { ++ x[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; ++ } ++ ++ ggml_vk_quantize_data(x, qx, x_ne, quant); ++ ++ for (size_t i = 0; i < y_ne; i++) { ++ // y[i] = rand() / (float)RAND_MAX; ++ y[i] = (i % k == i / k) ? 1.0f : 0.0f; ++ } ++ ++ ggml_pipeline_request_descriptor_sets(ctx->device, p, num_it); ++ if (split_k > 1) { ++ ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, num_it); ++ ++ if (ctx->prealloc_split_k == nullptr || ctx->prealloc_split_k->size < sizeof(float) * d_ne * split_k) { ++ // Resize buffer ++ if (ctx->prealloc_split_k != nullptr) { ++ ggml_vk_destroy_buffer(ctx->prealloc_split_k); ++ } ++ ctx->prealloc_split_k = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne * split_k, vk::MemoryPropertyFlagBits::eDeviceLocal); ++ } ++ } ++ ++ ggml_pipeline_allocate_descriptor_sets(ctx->device); ++ ++ ggml_vk_buffer_write(qx_buf, 0, qx, qx_sz); ++ ggml_vk_buffer_write(y_buf, 0, y, y_sz); ++ ++ vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); ++ ggml_vk_ctx_begin(ctx->device, subctx); ++ for (size_t i = 0; i < num_it; i++) { ++ ggml_vk_matmul( ++ ctx, subctx, p, ggml_vk_subbuffer(qx_buf), ggml_vk_subbuffer(y_buf), ggml_vk_subbuffer(d_buf), ggml_vk_subbuffer(ctx->prealloc_split_k), ++ m, n, k, ++ k, k, m, k*m, k*n, m*n, ++ split_k, batch, batch, batch, 1, 1 ++ ); ++ } ++ ggml_vk_ctx_end(subctx); ++ ++ auto begin = std::chrono::high_resolution_clock::now(); ++ ++ ggml_vk_submit(subctx, ctx->fence); ++ VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences"); ++ ctx->device->device.resetFences({ ctx->fence }); ++ ++ auto end = std::chrono::high_resolution_clock::now(); ++ ++ double time_ms = std::chrono::duration_cast(end-begin).count() / 1000.0; ++ ggml_vk_buffer_read(d_buf, 0, d, d_sz); ++ ++ ggml_init_params iparams = { ++ /*.mem_size =*/ 1024*1024*1024, ++ /*.mem_buffer =*/ NULL, ++ /*.no_alloc =*/ true, ++ }; ++ ++ ggml_context * ggml_ctx = ggml_init(iparams); ++ ++ ggml_tensor * src0_ggml = ggml_new_tensor_3d(ggml_ctx, quant, k, m, batch); ++ ggml_tensor * src1_ggml = ggml_new_tensor_3d(ggml_ctx, GGML_TYPE_F32, k, n, batch); ++ ggml_tensor * tensor_ggml = ggml_mul_mat(ggml_ctx, src0_ggml, src1_ggml); ++ ++ src0_ggml->data = qx; ++ src1_ggml->data = y; ++ tensor_ggml->data = d_chk; ++ ++ ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); ++ ggml_build_forward_expand(cgraph, tensor_ggml); ++ ++ ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 1); ++ ++ ggml_free(ggml_ctx); ++ ++ double avg_err = 0.0; ++ int first_err_n = -1; ++ int first_err_m = -1; ++ int first_err_b = -1; ++ ++ for (size_t i = 0; i < m*n*batch; i++) { ++ double err = std::fabs(d[i] - d_chk[i]); ++ avg_err += err; ++ ++ if ((err > 0.05f || std::isnan(err)) && first_err_n == -1) { ++ first_err_b = i / (m * n); ++ first_err_n = (i % (m * n)) / m; ++ first_err_m = (i % (m * n)) % m; ++ } ++ } ++ ++ avg_err /= m * n; ++ ++ double tflops = 2.0*m*n*k*batch*num_it / (time_ms / 1000.0) / (1000.0*1000.0*1000.0*1000.0); ++ ++ std::cerr << "TEST MMQ " << shname << " m=" << m << " n=" << n << " k=" << k << " batch=" << batch << " split_k=" << split_k << " matmul " << time_ms / num_it << "ms " << tflops << " TFLOPS avg_err=" << avg_err << std::endl; ++ ++ if (avg_err > 0.01 || std::isnan(avg_err)) { ++ std::cerr << "m = " << first_err_m << " n = " << first_err_n << " b = " << first_err_b << std::endl; ++ std::cerr << "Actual result: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(d, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ std::cerr << std::endl; ++ std::cerr << "Expected result: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(d_chk, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ ++ if (split_k > 1) { ++ float * split_k_buf = (float *) malloc(sizeof(float) * d_ne * split_k); ++ ggml_vk_buffer_read(ctx->prealloc_split_k, 0, split_k_buf, sizeof(float) * d_ne * split_k); ++ ++ std::cerr << "d_buf0: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(split_k_buf, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ ++ std::cerr << "d_buf1: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(split_k_buf + d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ ++ std::cerr << "d_buf2: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(split_k_buf + 2 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ ++ std::cerr << "d_buf3: " << std::endl << std::endl; ++ ggml_vk_print_matrix_area(split_k_buf + 3 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); ++ ++ free(split_k_buf); ++ } ++ } ++ ++ ggml_vk_destroy_buffer(qx_buf); ++ ggml_vk_destroy_buffer(y_buf); ++ ggml_vk_destroy_buffer(d_buf); ++ ++ free(x); ++ free(qx); ++ free(y); ++ free(d); ++ free(d_chk); ++} ++#endif ++ ++static void ggml_vk_preallocate_buffers(ggml_backend_vk_context * ctx) { ++#if defined(GGML_VULKAN_RUN_TESTS) ++ const std::vector vals { ++ 512, 512, 128, ++ 128, 512, 512, ++ 4096, 512, 4096, ++ 11008, 512, 4096, ++ 4096, 512, 11008, ++ 32000, 512, 4096, ++ 8, 8, 8, ++ 100, 46, 576, ++ 623, 111, 128, ++ 100, 46, 558, ++ 512, 1, 256, ++ 128, 110, 622, ++ 511, 511, 127, ++ 511, 511, 7, ++ 511, 511, 17, ++ 49, 49, 128, ++ 128, 49, 49, ++ 4096, 49, 4096, ++ }; ++ const size_t num_it = 100; ++ ++ for (size_t i = 0; i < vals.size(); i += 3) { ++ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0); ++ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1); ++ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2); ++ std::cerr << '\n'; ++ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0); ++ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1); ++ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2); ++ std::cerr << '\n'; ++ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0); ++ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1); ++ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2); ++ std::cerr << '\n' << std::endl; ++ ++ if (vals[i + 2] % 32 == 0) { ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0, GGML_TYPE_Q4_0); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1, GGML_TYPE_Q4_0); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2, GGML_TYPE_Q4_0); ++ std::cerr << '\n'; ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0, GGML_TYPE_Q4_0); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1, GGML_TYPE_Q4_0); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2, GGML_TYPE_Q4_0); ++ std::cerr << '\n'; ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0, GGML_TYPE_Q4_0); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1, GGML_TYPE_Q4_0); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2, GGML_TYPE_Q4_0); ++ std::cerr << '\n' << std::endl; ++ } ++ ++ if (vals[i + 2] % 256 == 0) { ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0, GGML_TYPE_Q4_K); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1, GGML_TYPE_Q4_K); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2, GGML_TYPE_Q4_K); ++ std::cerr << '\n'; ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0, GGML_TYPE_Q4_K); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1, GGML_TYPE_Q4_K); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2, GGML_TYPE_Q4_K); ++ std::cerr << '\n'; ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0, GGML_TYPE_Q4_K); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1, GGML_TYPE_Q4_K); ++ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2, GGML_TYPE_Q4_K); ++ std::cerr << '\n' << std::endl; ++ } ++ } ++ ++ GGML_ABORT("fatal error"); ++#endif ++ ++ if (ctx->prealloc_x == nullptr || (ctx->prealloc_size_x > 0 && ctx->prealloc_x->size < ctx->prealloc_size_x)) { ++ VK_LOG_MEMORY("ggml_vk_preallocate_buffers(x_size: " << ctx->prealloc_size_x << ")"); ++ // Resize buffer ++ if (ctx->prealloc_x != nullptr) { ++ ggml_vk_destroy_buffer(ctx->prealloc_x); ++ } ++ ctx->prealloc_x = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_x); ++ } ++ if (ctx->prealloc_y == nullptr || (ctx->prealloc_size_y > 0 && ctx->prealloc_y->size < ctx->prealloc_size_y)) { ++ VK_LOG_MEMORY("ggml_vk_preallocate_buffers(y_size: " << ctx->prealloc_size_y << ")"); ++ // Resize buffer ++ if (ctx->prealloc_y != nullptr) { ++ ggml_vk_destroy_buffer(ctx->prealloc_y); ++ } ++ ctx->prealloc_y = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_y); ++ } ++ if (ctx->prealloc_split_k == nullptr || (ctx->prealloc_size_split_k > 0 && ctx->prealloc_split_k->size < ctx->prealloc_size_split_k)) { ++ VK_LOG_MEMORY("ggml_vk_preallocate_buffers(split_k_size: " << ctx->prealloc_size_split_k << ")"); ++ // Resize buffer ++ if (ctx->prealloc_split_k != nullptr) { ++ ggml_vk_destroy_buffer(ctx->prealloc_split_k); ++ } ++ ctx->prealloc_split_k = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_split_k); ++ } ++} ++ ++static bool ggml_vk_compute_forward(ggml_backend_vk_context* ctx, ggml_tensor* tensor, int tensor_idx, bool use_fence); ++ ++// Returns true if node has enqueued work into the queue, false otherwise ++// If submit is true the current all operations queued so far are being submitted to Vulkan to overlap cmdlist creation and GPU execution. ++static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_tensor * node, int node_idx, ggml_tensor *node_begin, int node_idx_begin, bool dryrun, bool last_node, bool submit){ ++ if (ggml_is_empty(node) || !node->buffer) { ++ return false; ++ } ++ ++ VK_LOG_DEBUG("ggml_vk_build_graph(" << node << ", " << ggml_op_name(node->op) << ")"); ++ ctx->semaphore_idx = 0; ++ ++ const ggml_tensor * src0 = node->src[0]; ++ const ggml_tensor * src1 = node->src[1]; ++ const ggml_tensor * src2 = node->src[2]; ++ const ggml_tensor * src3 = node->src[3]; ++ ++ switch (node->op) { ++ // Return on empty ops to avoid generating a compute_ctx and setting exit_tensor ++ case GGML_OP_RESHAPE: ++ case GGML_OP_VIEW: ++ case GGML_OP_PERMUTE: ++ case GGML_OP_TRANSPOSE: ++ case GGML_OP_NONE: ++ return false; ++ case GGML_OP_UNARY: ++ switch (ggml_get_unary_op(node)) { ++ case GGML_UNARY_OP_SILU: ++ case GGML_UNARY_OP_GELU: ++ case GGML_UNARY_OP_GELU_QUICK: ++ case GGML_UNARY_OP_RELU: ++ case GGML_UNARY_OP_TANH: ++ break; ++ default: ++ return false; ++ } ++ break; ++ case GGML_OP_REPEAT: ++ case GGML_OP_GET_ROWS: ++ case GGML_OP_ADD: ++ case GGML_OP_ACC: ++ case GGML_OP_MUL: ++ case GGML_OP_DIV: ++ case GGML_OP_CONCAT: ++ case GGML_OP_UPSCALE: ++ case GGML_OP_SCALE: ++ case GGML_OP_SQR: ++ case GGML_OP_SIN: ++ case GGML_OP_COS: ++ case GGML_OP_CLAMP: ++ case GGML_OP_PAD: ++ case GGML_OP_CPY: ++ case GGML_OP_CONT: ++ case GGML_OP_DUP: ++ case GGML_OP_NORM: ++ case GGML_OP_GROUP_NORM: ++ case GGML_OP_RMS_NORM: ++ case GGML_OP_DIAG_MASK_INF: ++ case GGML_OP_SOFT_MAX: ++ case GGML_OP_ROPE: ++ case GGML_OP_MUL_MAT: ++ case GGML_OP_MUL_MAT_ID: ++ case GGML_OP_ARGSORT: ++ case GGML_OP_SUM_ROWS: ++ case GGML_OP_IM2COL: ++ case GGML_OP_TIMESTEP_EMBEDDING: ++ case GGML_OP_POOL_2D: ++ case GGML_OP_RWKV_WKV6: ++ case GGML_OP_LEAKY_RELU: ++ case GGML_OP_FLASH_ATTN_EXT: ++ break; ++ default: ++ std::cerr << "ggml_vulkan: Error: Missing op: " << ggml_op_name(node->op) << std::endl; ++ GGML_ABORT("fatal error"); ++ return false; ++ } ++ ++ vk_context compute_ctx; ++ ++ if (!dryrun) { ++ if (ctx->compute_ctx.expired()) { ++ compute_ctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); ++ ctx->compute_ctx = compute_ctx; ++ ggml_vk_ctx_begin(ctx->device, compute_ctx); ++ } else { ++ compute_ctx = ctx->compute_ctx.lock(); ++ } ++ } else { ++ switch (node->op) { ++ case GGML_OP_REPEAT: ++ case GGML_OP_ACC: ++ case GGML_OP_GET_ROWS: ++ case GGML_OP_ADD: ++ case GGML_OP_MUL: ++ case GGML_OP_DIV: ++ case GGML_OP_CONCAT: ++ case GGML_OP_UPSCALE: ++ case GGML_OP_SCALE: ++ case GGML_OP_SQR: ++ case GGML_OP_SIN: ++ case GGML_OP_COS: ++ case GGML_OP_CLAMP: ++ case GGML_OP_PAD: ++ case GGML_OP_CPY: ++ case GGML_OP_CONT: ++ case GGML_OP_DUP: ++ case GGML_OP_NORM: ++ case GGML_OP_GROUP_NORM: ++ case GGML_OP_RMS_NORM: ++ case GGML_OP_UNARY: ++ case GGML_OP_DIAG_MASK_INF: ++ case GGML_OP_SOFT_MAX: ++ case GGML_OP_ROPE: ++ case GGML_OP_ARGSORT: ++ case GGML_OP_SUM_ROWS: ++ case GGML_OP_IM2COL: ++ case GGML_OP_TIMESTEP_EMBEDDING: ++ case GGML_OP_POOL_2D: ++ case GGML_OP_LEAKY_RELU: ++ { ++ // These operations all go through ggml_vk_op_f32, so short-circuit and ++ // do the only thing needed for the dryrun. ++ vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, src2, node, node->op); ++ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); ++ return false; ++ } ++ default: ++ break; ++ } ++ } ++ ++ switch (node->op) { ++ case GGML_OP_REPEAT: ++ ggml_vk_repeat(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_ACC: ++ ggml_vk_acc(ctx, compute_ctx, src0, src1, node, dryrun); ++ ++ break; ++ case GGML_OP_GET_ROWS: ++ ggml_vk_get_rows(ctx, compute_ctx, src0, src1, node, dryrun); ++ ++ break; ++ case GGML_OP_ADD: ++ ggml_vk_add(ctx, compute_ctx, src0, src1, node, dryrun); ++ ++ break; ++ case GGML_OP_MUL: ++ ggml_vk_mul(ctx, compute_ctx, src0, src1, node, dryrun); ++ ++ break; ++ case GGML_OP_DIV: ++ ggml_vk_div(ctx, compute_ctx, src0, src1, node, dryrun); ++ ++ break; ++ case GGML_OP_CONCAT: ++ ggml_vk_concat(ctx, compute_ctx, src0, src1, node, dryrun); ++ ++ break; ++ case GGML_OP_UPSCALE: ++ ggml_vk_upscale(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_SCALE: ++ ggml_vk_scale(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_SQR: ++ ggml_vk_sqr(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_SIN: ++ ggml_vk_sin(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_COS: ++ ggml_vk_cos(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_CLAMP: ++ ggml_vk_clamp(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_PAD: ++ ggml_vk_pad(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_CPY: ++ case GGML_OP_CONT: ++ case GGML_OP_DUP: ++ ggml_vk_cpy(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_NORM: ++ ggml_vk_norm(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_GROUP_NORM: ++ ggml_vk_group_norm(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_RMS_NORM: ++ ggml_vk_rms_norm(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_UNARY: ++ switch (ggml_get_unary_op(node)) { ++ case GGML_UNARY_OP_SILU: ++ case GGML_UNARY_OP_GELU: ++ case GGML_UNARY_OP_GELU_QUICK: ++ case GGML_UNARY_OP_RELU: ++ case GGML_UNARY_OP_TANH: ++ ggml_vk_unary(ctx, compute_ctx, src0, node, dryrun); ++ break; ++ default: ++ return false; ++ } ++ break; ++ case GGML_OP_DIAG_MASK_INF: ++ ggml_vk_diag_mask_inf(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_SOFT_MAX: ++ ggml_vk_soft_max(ctx, compute_ctx, src0, src1, node, dryrun); ++ ++ break; ++ case GGML_OP_ROPE: ++ ggml_vk_rope(ctx, compute_ctx, src0, src1, src2, node, dryrun); ++ ++ break; ++ case GGML_OP_ARGSORT: ++ ggml_vk_argsort(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_SUM_ROWS: ++ ggml_vk_sum_rows(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_IM2COL: ++ ggml_vk_im2col(ctx, compute_ctx, src0, src1, node, dryrun); ++ ++ break; ++ case GGML_OP_TIMESTEP_EMBEDDING: ++ ggml_vk_timestep_embedding(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_POOL_2D: ++ ggml_vk_pool_2d(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_LEAKY_RELU: ++ ggml_vk_leaky_relu(ctx, compute_ctx, src0, node, dryrun); ++ ++ break; ++ case GGML_OP_MUL_MAT: ++ ggml_vk_mul_mat(ctx, compute_ctx, src0, src1, node, dryrun); ++ ++ break; ++ case GGML_OP_MUL_MAT_ID: ++ ggml_vk_mul_mat_id(ctx, compute_ctx, src0, src1, src2, node, dryrun); ++ ++ break; ++ ++ case GGML_OP_FLASH_ATTN_EXT: ++ ggml_vk_flash_attn(ctx, compute_ctx, src0, src1, src2, src3, node, dryrun); ++ ++ break; ++ ++ case GGML_OP_RWKV_WKV6: ++ ggml_vk_rwkv_wkv6(ctx, compute_ctx, node, dryrun); ++ ++ break; ++ default: ++ return false; ++ } ++ ++ if (dryrun) { ++ return false; ++ } ++ ++ ctx->tensor_ctxs[node_idx] = compute_ctx; ++ ++#if defined(GGML_VULKAN_CHECK_RESULTS) || defined(GGML_VULKAN_PERF) ++ // Force context reset on each node so that each tensor ends up in its own context ++ // and can be run and compared to its CPU equivalent separately ++ last_node = true; ++#endif ++ ++ if (submit || last_node) { ++ ggml_vk_ctx_end(compute_ctx); ++ ++ // TODO probably it'd be better to pass a exit_node flag to ggml_vk_compute_forward ++ if (last_node) { ++ compute_ctx->exit_tensor_idx = node_idx_begin; ++ } ++ else { ++ compute_ctx->exit_tensor_idx = -1; ++ } ++ ++ ctx->compute_ctx.reset(); ++ ++ bool ok = ggml_vk_compute_forward(ctx, node_begin, node_idx_begin, false); ++ if (!ok) { ++ if (node->op == GGML_OP_UNARY) { ++ std::cerr << __func__ << ": error: op not supported UNARY " << node->name << " (" << ggml_unary_op_name(static_cast(node->op_params[0])) << ")" << std::endl; ++ } ++ else { ++ std::cerr << __func__ << ": error: op not supported " << node->name << " (" << ggml_op_name(node->op) << ")" << std::endl; ++ } ++ } ++ ++ } ++ return true; ++} ++ ++static bool ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_tensor * tensor, int tensor_idx, bool use_fence = true){ ++ ggml_backend_buffer * buf = nullptr; ++ ++ switch (tensor->op) { ++ case GGML_OP_ADD: ++ case GGML_OP_ACC: ++ case GGML_OP_GET_ROWS: ++ case GGML_OP_MUL: ++ case GGML_OP_DIV: ++ case GGML_OP_CONCAT: ++ case GGML_OP_UPSCALE: ++ case GGML_OP_SCALE: ++ case GGML_OP_SQR: ++ case GGML_OP_SIN: ++ case GGML_OP_COS: ++ case GGML_OP_CLAMP: ++ case GGML_OP_PAD: ++ case GGML_OP_CPY: ++ case GGML_OP_CONT: ++ case GGML_OP_DUP: ++ case GGML_OP_NORM: ++ case GGML_OP_GROUP_NORM: ++ case GGML_OP_RMS_NORM: ++ case GGML_OP_DIAG_MASK_INF: ++ case GGML_OP_SOFT_MAX: ++ case GGML_OP_ROPE: ++ case GGML_OP_RESHAPE: ++ case GGML_OP_VIEW: ++ case GGML_OP_PERMUTE: ++ case GGML_OP_TRANSPOSE: ++ case GGML_OP_NONE: ++ case GGML_OP_ARGSORT: ++ case GGML_OP_SUM_ROWS: ++ case GGML_OP_IM2COL: ++ case GGML_OP_TIMESTEP_EMBEDDING: ++ case GGML_OP_POOL_2D: ++ case GGML_OP_RWKV_WKV6: ++ case GGML_OP_LEAKY_RELU: ++ case GGML_OP_REPEAT: ++ buf = tensor->buffer; ++ ++ break; ++ case GGML_OP_UNARY: ++ switch (ggml_get_unary_op(tensor)) { ++ case GGML_UNARY_OP_SILU: ++ case GGML_UNARY_OP_GELU: ++ case GGML_UNARY_OP_GELU_QUICK: ++ case GGML_UNARY_OP_RELU: ++ case GGML_UNARY_OP_TANH: ++ buf = tensor->buffer; ++ break; ++ default: ++ return false; ++ } ++ break; ++ case GGML_OP_MUL_MAT: ++ case GGML_OP_MUL_MAT_ID: ++ case GGML_OP_FLASH_ATTN_EXT: ++ buf = tensor->buffer; ++ ++ break; ++ default: ++ return false; ++ } ++ ++ if (buf == nullptr) { ++ return false; ++ } ++ ++ VK_LOG_DEBUG("ggml_vk_compute_forward(" << tensor << ", name=" << tensor->name << ", op=" << ggml_op_name(tensor->op) << ", type=" << tensor->type << ", ne0=" << tensor->ne[0] << ", ne1=" << tensor->ne[1] << ", ne2=" << tensor->ne[2] << ", ne3=" << tensor->ne[3] << ", nb0=" << tensor->nb[0] << ", nb1=" << tensor->nb[1] << ", nb2=" << tensor->nb[2] << ", nb3=" << tensor->nb[3] << ", view_src=" << tensor->view_src << ", view_offs=" << tensor->view_offs << ")"); ++ ++ vk_context subctx = ctx->tensor_ctxs[tensor_idx].lock(); ++ ++ // always wait for the GPU work to be done for the last submit ++ if (tensor_idx == subctx->exit_tensor_idx) { ++ use_fence = true; ++ } ++ ++ // Only run if ctx hasn't been submitted yet ++ if (!subctx->seqs.empty()) { ++#ifdef GGML_VULKAN_CHECK_RESULTS ++ ggml_vk_check_results_0(tensor); ++ use_fence = true; ++#endif ++ ++ // Do staging buffer copies ++ for (auto& cpy : subctx->in_memcpys) { ++ memcpy(cpy.dst, cpy.src, cpy.n); ++ } ++ ++ ggml_vk_submit(subctx, use_fence ? ctx->fence : vk::Fence{}); ++ ++ if (use_fence) { ++ VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_compute_forward waitForFences"); ++ ++ ctx->device->device.resetFences({ ctx->fence }); ++ } ++#ifdef GGML_VULKAN_CHECK_RESULTS ++ ggml_vk_check_results_1(tensor); ++#endif ++ } ++ ++ if (tensor_idx == subctx->exit_tensor_idx) { ++ // Do staging buffer copies ++ for (auto& cpy : subctx->out_memcpys) { ++ memcpy(cpy.dst, cpy.src, cpy.n); ++ } ++ subctx->in_memcpys.clear(); ++ subctx->out_memcpys.clear(); ++ } ++ ++ return true; ++} ++ ++// Clean up after graph processing is done ++static void ggml_vk_graph_cleanup(ggml_backend_vk_context * ctx) { ++ VK_LOG_DEBUG("ggml_vk_graph_cleanup()"); ++ for (auto& buffer : ctx->gc.temp_buffers) { ++ ggml_vk_pool_free(ctx, buffer); ++ } ++ ctx->gc.temp_buffers.clear(); ++ ++ for (auto& dsr : ctx->device->pipeline_descriptor_set_requirements) { ++ vk_pipeline_ref plr = ctx->device->pipelines[dsr.first]; ++ ++ if (plr.expired()) { ++ continue; ++ } ++ ++ vk_pipeline pl = plr.lock(); ++ ggml_pipeline_cleanup(pl); ++ } ++ ++ ggml_vk_queue_cleanup(ctx->device, ctx->device->compute_queue); ++ ggml_vk_queue_cleanup(ctx->device, ctx->device->transfer_queue); ++ ++ for (size_t i = 0; i < ctx->gc.semaphores.size(); i++) { ++ ctx->device->device.destroySemaphore({ ctx->gc.semaphores[i].s }); ++ } ++ ctx->gc.semaphores.clear(); ++ ++ for (size_t i = 0; i < ctx->gc.tl_semaphores.size(); i++) { ++ ctx->device->device.destroySemaphore({ ctx->gc.tl_semaphores[i].s }); ++ } ++ ctx->gc.tl_semaphores.clear(); ++ ctx->semaphore_idx = 0; ++ ++ ctx->event_idx = 0; ++ ++ for (auto& event : ctx->gc.events) { ++ ctx->device->device.resetEvent(event); ++ } ++ ++ ctx->tensor_ctxs.clear(); ++ ctx->gc.contexts.clear(); ++ ctx->device->pipeline_descriptor_set_requirements.clear(); ++} ++ ++// Clean up on backend free ++static void ggml_vk_cleanup(ggml_backend_vk_context * ctx) { ++ VK_LOG_DEBUG("ggml_vk_cleanup(" << ctx->name << ")"); ++ ggml_vk_graph_cleanup(ctx); ++ ++ ggml_vk_destroy_buffer(ctx->prealloc_x); ++ ggml_vk_destroy_buffer(ctx->prealloc_y); ++ ggml_vk_destroy_buffer(ctx->prealloc_split_k); ++ ++ for (auto& buffer : ctx->buffer_pool) { ++ ggml_vk_destroy_buffer(buffer); ++ } ++ ++ ctx->prealloc_size_x = 0; ++ ctx->prealloc_size_y = 0; ++ ctx->prealloc_size_split_k = 0; ++ ++ for (auto& event : ctx->gc.events) { ++ ctx->device->device.destroyEvent(event); ++ } ++ ctx->gc.events.clear(); ++ ++ ctx->device->device.destroyFence(ctx->fence); ++} ++ ++static int ggml_vk_get_device_count() { ++ ggml_vk_instance_init(); ++ ++ return vk_instance.device_indices.size(); ++} ++ ++static void ggml_vk_get_device_description(int device, char * description, size_t description_size) { ++ ggml_vk_instance_init(); ++ ++ std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); ++ ++ vk::PhysicalDeviceProperties props; ++ devices[device].getProperties(&props); ++ ++ snprintf(description, description_size, "%s", props.deviceName.data()); ++} ++ ++// backend interface ++ ++#define UNUSED GGML_UNUSED ++ ++// device backend ++ ++static bool ggml_backend_buffer_is_vk(ggml_backend_buffer_t buffer) { ++ return buffer->buft->iface.get_name == ggml_backend_vk_buffer_type_name; ++} ++ ++static void ggml_backend_vk_buffer_free_buffer(ggml_backend_buffer_t buffer) { ++ VK_LOG_MEMORY("ggml_backend_vk_buffer_free_buffer()"); ++ ggml_backend_vk_buffer_context * ctx = (ggml_backend_vk_buffer_context *)buffer->context; ++ ggml_vk_destroy_buffer(ctx->dev_buffer); ++ delete ctx; ++} ++ ++static void * ggml_backend_vk_buffer_get_base(ggml_backend_buffer_t buffer) { ++ return vk_ptr_base; ++ ++ UNUSED(buffer); ++} ++ ++static void ggml_backend_vk_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { ++ VK_LOG_DEBUG("ggml_backend_vk_buffer_init_tensor(" << buffer << " (" << buffer->context << "), " << tensor << ")"); ++ if (tensor->view_src != nullptr) { ++ GGML_ASSERT(tensor->view_src->buffer->buft == buffer->buft); ++ } ++} ++ ++static void ggml_backend_vk_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { ++ VK_LOG_DEBUG("ggml_backend_vk_buffer_set_tensor(" << buffer << ", " << tensor << ", " << data << ", " << offset << ", " << size << ")"); ++ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)buffer->context; ++ vk_buffer buf = buf_ctx->dev_buffer; ++ ++ ggml_vk_buffer_write(buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); ++} ++ ++static void ggml_backend_vk_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { ++ VK_LOG_DEBUG("ggml_backend_vk_buffer_get_tensor(" << buffer << ", " << tensor << ", " << data << ", " << offset << ", " << size << ")"); ++ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)buffer->context; ++ ++ vk_buffer buf = buf_ctx->dev_buffer; ++ ++ ggml_vk_buffer_read(buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); ++} ++ ++static bool ggml_backend_vk_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) { ++ if (ggml_backend_buffer_is_vk(src->buffer)) { ++ ggml_backend_vk_buffer_context * src_buf_ctx = (ggml_backend_vk_buffer_context *)src->buffer->context; ++ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ++ ++ vk_buffer src_buf = src_buf_ctx->dev_buffer; ++ vk_buffer dst_buf = dst_buf_ctx->dev_buffer; ++ ++ ggml_vk_buffer_copy(dst_buf, vk_tensor_offset(dst) + dst->view_offs, src_buf, vk_tensor_offset(src) + src->view_offs, ggml_nbytes(src)); ++ ++ return true; ++ } ++ return false; ++ ++ UNUSED(buffer); ++} ++ ++static void ggml_backend_vk_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { ++ ggml_backend_vk_buffer_context * ctx = (ggml_backend_vk_buffer_context *)buffer->context; ++ ++ ggml_vk_buffer_memset(ctx->dev_buffer, 0, value, buffer->size); ++} ++ ++static ggml_backend_buffer_i ggml_backend_vk_buffer_interface = { ++ /* .free_buffer = */ ggml_backend_vk_buffer_free_buffer, ++ /* .get_base = */ ggml_backend_vk_buffer_get_base, ++ /* .init_tensor = */ ggml_backend_vk_buffer_init_tensor, ++ /* .memset_tensor = */ NULL, ++ /* .set_tensor = */ ggml_backend_vk_buffer_set_tensor, ++ /* .get_tensor = */ ggml_backend_vk_buffer_get_tensor, ++ /* .cpy_tensor = */ ggml_backend_vk_buffer_cpy_tensor, ++ /* .clear = */ ggml_backend_vk_buffer_clear, ++ /* .reset = */ NULL, ++}; ++ ++// vk buffer type ++static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft) { ++ ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *)buft->context; ++ ++ return ctx->name.c_str(); ++} ++ ++static ggml_backend_buffer_t ggml_backend_vk_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { ++ VK_LOG_MEMORY("ggml_backend_vk_buffer_type_alloc_buffer(" << size << ")"); ++ ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; ++ ++ vk_buffer dev_buffer = nullptr; ++ try { ++ dev_buffer = ggml_vk_create_buffer_device(ctx->device, size); ++ } catch (const vk::SystemError& e) { ++ return nullptr; ++ } ++ ++ ggml_backend_vk_buffer_context * bufctx = new ggml_backend_vk_buffer_context(ctx->device, std::move(dev_buffer), ctx->name); ++ ++ return ggml_backend_buffer_init(buft, ggml_backend_vk_buffer_interface, bufctx, size); ++} ++ ++static size_t ggml_backend_vk_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { ++ ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; ++ return ctx->device->properties.limits.minStorageBufferOffsetAlignment; ++} ++ ++static size_t ggml_backend_vk_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { ++ ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; ++ return ctx->device->max_memory_allocation_size; ++} ++ ++static size_t ggml_backend_vk_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { ++ return ggml_nbytes(tensor); ++ ++ UNUSED(buft); ++} ++ ++ggml_backend_buffer_type_t ggml_backend_vk_buffer_type(size_t dev_num) { ++ ggml_vk_instance_init(); ++ ++ VK_LOG_DEBUG("ggml_backend_vk_buffer_type(" << dev_num << ")"); ++ ++ vk_device dev = ggml_vk_get_device(dev_num); ++ ++ return &dev->buffer_type; ++} ++ ++// host buffer type ++ ++static const char * ggml_backend_vk_host_buffer_type_name(ggml_backend_buffer_type_t buft) { ++ return GGML_VK_NAME "_Host"; ++ ++ UNUSED(buft); ++} ++ ++static const char * ggml_backend_vk_host_buffer_name(ggml_backend_buffer_t buffer) { ++ return GGML_VK_NAME "_Host"; ++ ++ UNUSED(buffer); ++} ++ ++static void ggml_backend_vk_host_buffer_free_buffer(ggml_backend_buffer_t buffer) { ++ VK_LOG_MEMORY("ggml_backend_vk_host_buffer_free_buffer()"); ++ ggml_vk_host_free(vk_instance.devices[0], buffer->context); ++} ++ ++static ggml_backend_buffer_t ggml_backend_vk_host_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { ++ VK_LOG_MEMORY("ggml_backend_vk_host_buffer_type_alloc_buffer(" << size << ")"); ++ ++ size += 32; // Behave like the CPU buffer type ++ void * ptr = nullptr; ++ try { ++ ptr = ggml_vk_host_malloc(vk_instance.devices[0], size); ++ } catch (vk::SystemError& e) { ++ std::cerr << "ggml_vulkan: Failed to allocate pinned memory." << std::endl; ++ std::cerr << "ggml_vulkan: " << e.what() << std::endl; ++ // fallback to cpu buffer ++ return ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); ++ } ++ ++ ggml_backend_buffer_t buffer = ggml_backend_cpu_buffer_from_ptr(ptr, size); ++ buffer->buft = buft; ++ buffer->iface.free_buffer = ggml_backend_vk_host_buffer_free_buffer; ++ ++ return buffer; ++ ++ UNUSED(buft); ++} ++ ++static size_t ggml_backend_vk_host_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { ++ return vk_instance.devices[0]->properties.limits.minMemoryMapAlignment; ++ ++ UNUSED(buft); ++} ++ ++// Should be changed to return device-specific host buffer type ++// but that probably requires changes in llama.cpp ++ggml_backend_buffer_type_t ggml_backend_vk_host_buffer_type() { ++ static struct ggml_backend_buffer_type ggml_backend_vk_buffer_type_host = { ++ /* .iface = */ { ++ /* .get_name = */ ggml_backend_vk_host_buffer_type_name, ++ /* .alloc_buffer = */ ggml_backend_vk_host_buffer_type_alloc_buffer, ++ /* .get_alignment = */ ggml_backend_vk_host_buffer_type_get_alignment, ++ /* .get_max_size = */ NULL, // defaults to SIZE_MAX ++ /* .get_alloc_size = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size, ++ /* .is_host = */ ggml_backend_cpu_buffer_type()->iface.is_host, ++ }, ++ /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), 0), ++ /* .context = */ nullptr, ++ }; ++ ++ // Make sure device 0 is initialized ++ ggml_vk_instance_init(); ++ ggml_vk_get_device(0); ++ ++ return &ggml_backend_vk_buffer_type_host; ++} ++ ++ ++// backend ++ ++static const char * ggml_backend_vk_name(ggml_backend_t backend) { ++ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; ++ ++ return ctx->name.c_str(); ++} ++ ++static void ggml_backend_vk_free(ggml_backend_t backend) { ++ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; ++ VK_LOG_DEBUG("ggml_backend_vk_free(" << ctx->name << ")"); ++ ++ ggml_vk_cleanup(ctx); ++ ++ delete ctx; ++ delete backend; ++} ++ ++static ggml_backend_buffer_type_t ggml_backend_vk_get_default_buffer_type(ggml_backend_t backend) { ++ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; ++ ++ return &ctx->device->buffer_type; ++} ++ ++static void ggml_backend_vk_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { ++ VK_LOG_DEBUG("ggml_backend_vk_set_tensor_async(" << size << ")"); ++ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; ++ GGML_ASSERT((tensor->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || tensor->buffer->buft == ggml_backend_vk_host_buffer_type()) && "unsupported buffer type"); ++ ++ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; ++ ++ vk_context transfer_ctx; ++ ++ if (ctx->transfer_ctx.expired()) { ++ // Initialize new transfer context ++ transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); ++ ctx->transfer_ctx = transfer_ctx; ++ ggml_vk_ctx_begin(ctx->device, transfer_ctx); ++ } else { ++ transfer_ctx = ctx->transfer_ctx.lock(); ++ } ++ ++ vk_buffer buf = buf_ctx->dev_buffer; ++ ++ ggml_vk_buffer_write_async(transfer_ctx, buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); ++} ++ ++static void ggml_backend_vk_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { ++ VK_LOG_DEBUG("ggml_backend_vk_get_tensor_async(" << size << ")"); ++ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; ++ GGML_ASSERT((tensor->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || tensor->buffer->buft == ggml_backend_vk_host_buffer_type()) && "unsupported buffer type"); ++ ++ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; ++ ++ vk_context transfer_ctx; ++ ++ if (ctx->transfer_ctx.expired()) { ++ // Initialize new transfer context ++ transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); ++ ctx->transfer_ctx = transfer_ctx; ++ ggml_vk_ctx_begin(ctx->device, transfer_ctx); ++ } else { ++ transfer_ctx = ctx->transfer_ctx.lock(); ++ } ++ ++ vk_buffer buf = buf_ctx->dev_buffer; ++ ++ ggml_vk_buffer_read_async(transfer_ctx, buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); ++} ++ ++static bool ggml_backend_vk_cpy_tensor_async(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst) { ++ VK_LOG_DEBUG("ggml_backend_vk_cpy_tensor_async()"); ++ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; ++ if ((dst->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || dst->buffer->buft == ggml_backend_vk_host_buffer_type()) && ggml_backend_buffer_is_vk(src->buffer)) { ++ ggml_backend_vk_buffer_context * src_buf_ctx = (ggml_backend_vk_buffer_context *)src->buffer->context; ++ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ++ ++ vk_context transfer_ctx; ++ ++ if (ctx->transfer_ctx.expired()) { ++ // Initialize new transfer context ++ transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); ++ ctx->transfer_ctx = transfer_ctx; ++ ggml_vk_ctx_begin(ctx->device, transfer_ctx); ++ } else { ++ transfer_ctx = ctx->transfer_ctx.lock(); ++ } ++ ++ vk_buffer src_buf = src_buf_ctx->dev_buffer; ++ vk_buffer dst_buf = dst_buf_ctx->dev_buffer; ++ ++ ggml_vk_buffer_copy_async(transfer_ctx, dst_buf, vk_tensor_offset(dst) + dst->view_offs, src_buf, vk_tensor_offset(src) + src->view_offs, ggml_nbytes(src)); ++ return true; ++ } ++ ++ return false; ++} ++ ++static void ggml_backend_vk_synchronize(ggml_backend_t backend) { ++ VK_LOG_DEBUG("ggml_backend_vk_synchronize()"); ++ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; ++ if(ctx->transfer_ctx.expired()) { ++ return; ++ } ++ ++ vk_context transfer_ctx = ctx->transfer_ctx.lock(); ++ ++ ggml_vk_ctx_end(transfer_ctx); ++ ++ for (auto& cpy : transfer_ctx->in_memcpys) { ++ memcpy(cpy.dst, cpy.src, cpy.n); ++ } ++ ++ ggml_vk_submit(transfer_ctx, ctx->fence); ++ VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_backend_vk_synchronize waitForFences"); ++ ctx->device->device.resetFences({ ctx->fence }); ++ ++ for (auto& cpy : transfer_ctx->out_memcpys) { ++ memcpy(cpy.dst, cpy.src, cpy.n); ++ } ++ ++ ctx->transfer_ctx.reset(); ++} ++ ++static bool ggml_vk_is_empty(ggml_tensor * node) { ++ return ggml_is_empty(node) || node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE; ++} ++ ++static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { ++ VK_LOG_DEBUG("ggml_backend_vk_graph_compute(" << cgraph->n_nodes << " nodes)"); ++ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; ++ ++ for (int i = 0; i < cgraph->n_nodes; i++) { ++ ggml_vk_build_graph(ctx, cgraph->nodes[i], i, nullptr, 0, true, false, false); ++ } ++ ggml_vk_preallocate_buffers(ctx); ++ ggml_pipeline_allocate_descriptor_sets(ctx->device); ++ ++ int last_node = cgraph->n_nodes - 1; ++ ++ // If the last op in the cgraph isn't backend GPU, the command buffer doesn't get closed properly ++ while (last_node > 0 && ggml_vk_is_empty(cgraph->nodes[last_node])) { ++ last_node -= 1; ++ } ++ ++ // Reserve tensor context space for all nodes ++ ctx->tensor_ctxs.resize(cgraph->n_nodes); ++ ++ bool first_node_in_batch = true; // true if next node will be first node in a batch ++ int submit_node_idx = 0; // index to first node in a batch ++ ++ // Submit work every nodes_per_submit nodes to overlap CPU cmdbuffer generation with GPU execution. ++ // Start with a smaller count to get work submitted right away, and increase it after each submit. ++ int nodes_per_submit = 20; ++ int submitted_nodes = 0; ++ int submit_count = 0; ++ for (int i = 0; i < cgraph->n_nodes; i++) { ++ if (first_node_in_batch) { ++ submit_node_idx = i; ++ } ++ ++ bool submit = (submitted_nodes >= nodes_per_submit) || (i == last_node); ++ ++ bool enqueued = ggml_vk_build_graph(ctx, cgraph->nodes[i], i, cgraph->nodes[submit_node_idx], submit_node_idx, false, i == last_node, submit); ++ ++ if (enqueued) { ++ ++submitted_nodes; ++ ++#ifndef GGML_VULKAN_CHECK_RESULTS ++ if (first_node_in_batch) { ++ first_node_in_batch = false; ++ } ++#endif ++ } ++ ++ if (submit) { ++ first_node_in_batch = true; ++ submitted_nodes = 0; ++ switch (submit_count) { ++ case 0: ++ nodes_per_submit = 50; ++ break; ++ default: ++ nodes_per_submit = 100; ++ break; ++ } ++ submit_count++; ++ } ++ } ++ ++#ifdef GGML_VULKAN_PERF ++ ctx->device->perf_logger->print_timings(); ++#endif ++ ++ ggml_vk_graph_cleanup(ctx); ++ ++ return GGML_STATUS_SUCCESS; ++ ++ UNUSED(backend); ++} ++ ++// TODO: enable async and synchronize ++static ggml_backend_i ggml_backend_vk_interface = { ++ /* .get_name = */ ggml_backend_vk_name, ++ /* .free = */ ggml_backend_vk_free, ++ /* .set_tensor_async = */ NULL, // ggml_backend_vk_set_tensor_async, ++ /* .get_tensor_async = */ NULL, // ggml_backend_vk_get_tensor_async, ++ /* .cpy_tensor_async = */ NULL, // ggml_backend_vk_cpy_tensor_async, ++ /* .synchronize = */ NULL, // ggml_backend_vk_synchronize, ++ /* .graph_plan_create = */ NULL, ++ /* .graph_plan_free = */ NULL, ++ /* .graph_plan_update = */ NULL, ++ /* .graph_plan_compute = */ NULL, ++ /* .graph_compute = */ ggml_backend_vk_graph_compute, ++ /* .event_record = */ NULL, ++ /* .event_wait = */ NULL, ++}; ++ ++static ggml_guid_t ggml_backend_vk_guid() { ++ static ggml_guid guid = { 0xb8, 0xf7, 0x4f, 0x86, 0x40, 0x3c, 0xe1, 0x02, 0x91, 0xc8, 0xdd, 0xe9, 0x02, 0x3f, 0xc0, 0x2b }; ++ return &guid; ++} ++ ++ggml_backend_t ggml_backend_vk_init(size_t dev_num) { ++ VK_LOG_DEBUG("ggml_backend_vk_init(" << dev_num << ")"); ++ ++ ggml_backend_vk_context * ctx = new ggml_backend_vk_context; ++ ggml_vk_init(ctx, dev_num); ++ ++ ggml_backend_t vk_backend = new ggml_backend { ++ /* .guid = */ ggml_backend_vk_guid(), ++ /* .interface = */ ggml_backend_vk_interface, ++ /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), dev_num), ++ /* .context = */ ctx, ++ }; ++ ++ return vk_backend; ++} ++ ++bool ggml_backend_is_vk(ggml_backend_t backend) { ++ return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_vk_guid()); ++} ++ ++int ggml_backend_vk_get_device_count() { ++ return ggml_vk_get_device_count(); ++} ++ ++void ggml_backend_vk_get_device_description(int device, char * description, size_t description_size) { ++ GGML_ASSERT(device < (int) vk_instance.device_indices.size()); ++ int dev_idx = vk_instance.device_indices[device]; ++ ggml_vk_get_device_description(dev_idx, description, description_size); ++} ++ ++void ggml_backend_vk_get_device_memory(int device, size_t * free, size_t * total) { ++ GGML_ASSERT(device < (int) vk_instance.device_indices.size()); ++ ++ vk::PhysicalDevice vkdev = vk_instance.instance.enumeratePhysicalDevices()[vk_instance.device_indices[device]]; ++ ++ vk::PhysicalDeviceMemoryProperties memprops = vkdev.getMemoryProperties(); ++ ++ for (const vk::MemoryHeap& heap : memprops.memoryHeaps) { ++ if (heap.flags & vk::MemoryHeapFlagBits::eDeviceLocal) { ++ *total = heap.size; ++ *free = heap.size; ++ break; ++ } ++ } ++} ++ ++////////////////////////// ++ ++struct ggml_backend_vk_device_context { ++ size_t device; ++ std::string name; ++ std::string description; ++}; ++ ++static const char * ggml_backend_vk_device_get_name(ggml_backend_dev_t dev) { ++ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; ++ return ctx->name.c_str(); ++} ++ ++static const char * ggml_backend_vk_device_get_description(ggml_backend_dev_t dev) { ++ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; ++ return ctx->description.c_str(); ++} ++ ++static void ggml_backend_vk_device_get_memory(ggml_backend_dev_t device, size_t * free, size_t * total) { ++ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)device->context; ++ ggml_backend_vk_get_device_memory(ctx->device, free, total); ++} ++ ++static ggml_backend_buffer_type_t ggml_backend_vk_device_get_buffer_type(ggml_backend_dev_t dev) { ++ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; ++ return ggml_backend_vk_buffer_type(ctx->device); ++} ++ ++static ggml_backend_buffer_type_t ggml_backend_vk_device_get_host_buffer_type(ggml_backend_dev_t dev) { ++ UNUSED(dev); ++ return ggml_backend_vk_host_buffer_type(); ++} ++ ++static enum ggml_backend_dev_type ggml_backend_vk_device_get_type(ggml_backend_dev_t dev) { ++ UNUSED(dev); ++ return GGML_BACKEND_DEVICE_TYPE_GPU; ++} ++ ++static void ggml_backend_vk_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props) { ++ props->name = ggml_backend_vk_device_get_name(dev); ++ props->description = ggml_backend_vk_device_get_description(dev); ++ props->type = ggml_backend_vk_device_get_type(dev); ++ ggml_backend_vk_device_get_memory(dev, &props->memory_free, &props->memory_total); ++ props->caps = { ++ /* .async = */ false, ++ /* .host_buffer = */ true, ++ /* .buffer_from_host_ptr = */ false, ++ /* .events = */ false, ++ }; ++} ++ ++static ggml_backend_t ggml_backend_vk_device_init(ggml_backend_dev_t dev, const char * params) { ++ UNUSED(params); ++ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; ++ return ggml_backend_vk_init(ctx->device); ++} ++ ++static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { ++ switch (op->op) { ++ case GGML_OP_UNARY: ++ switch (ggml_get_unary_op(op)) { ++ case GGML_UNARY_OP_GELU: ++ case GGML_UNARY_OP_GELU_QUICK: ++ case GGML_UNARY_OP_SILU: ++ case GGML_UNARY_OP_RELU: ++ case GGML_UNARY_OP_TANH: ++ return ggml_is_contiguous(op->src[0]); ++ default: ++ return false; ++ } ++ break; ++ case GGML_OP_MUL_MAT: ++ case GGML_OP_MUL_MAT_ID: ++ { ++ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; ++ const vk_device& device = ggml_vk_get_device(ctx->device); ++ if (op->op == GGML_OP_MUL_MAT_ID && !device->mul_mat_id_s && !device->mul_mat_id_m && !device->mul_mat_id_l) { ++ // If there's not enough shared memory for row_ids and the result tile, fallback to CPU ++ return false; ++ } ++ switch (op->src[0]->type) { ++ case GGML_TYPE_F32: ++ case GGML_TYPE_F16: ++ case GGML_TYPE_Q4_0: ++ case GGML_TYPE_Q4_1: ++ case GGML_TYPE_Q5_0: ++ case GGML_TYPE_Q5_1: ++ case GGML_TYPE_Q8_0: ++ case GGML_TYPE_Q2_K: ++ case GGML_TYPE_Q3_K: ++ case GGML_TYPE_Q4_K: ++ case GGML_TYPE_Q5_K: ++ case GGML_TYPE_Q6_K: ++ case GGML_TYPE_IQ4_NL: ++ break; ++ default: ++ return false; ++ } ++ struct ggml_tensor * a; ++ struct ggml_tensor * b; ++ if (op->op == GGML_OP_MUL_MAT) { ++ a = op->src[0]; ++ b = op->src[1]; ++ } else { ++ a = op->src[2]; ++ b = op->src[1]; ++ } ++ if (a->ne[3] != b->ne[3]) { ++ return false; ++ } ++ if (!(ggml_vk_dim01_contiguous(op->src[0]) || op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) || ++ !(ggml_vk_dim01_contiguous(op->src[1]) || op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16)) { ++ return false; ++ } ++ ++ return true; ++ } break; ++ case GGML_OP_FLASH_ATTN_EXT: ++ { ++ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; ++ if (!ggml_vk_get_device(ctx->device)->coopmat2) { ++ return false; ++ } ++ switch (op->src[0]->ne[0]) { ++ case 64: ++ case 80: ++ case 96: ++ case 112: ++ case 128: ++ case 256: ++ break; ++ default: ++ return false; ++ } ++ if (op->src[0]->type != GGML_TYPE_F32) { ++ return false; ++ } ++ if (op->type != GGML_TYPE_F32) { ++ return false; ++ } ++ if (op->src[3] && op->src[3]->type != GGML_TYPE_F16) { ++ return false; ++ } ++ // It's straightforward to support different K/V dequant, but would ++ // significantly increase the number of pipelines ++ if (op->src[1]->type != op->src[2]->type) { ++ return false; ++ } ++ switch (op->src[1]->type) { ++ case GGML_TYPE_F16: ++ case GGML_TYPE_Q4_0: ++ case GGML_TYPE_Q4_1: ++ case GGML_TYPE_Q5_0: ++ case GGML_TYPE_Q5_1: ++ case GGML_TYPE_Q8_0: ++ // K dequants currently disabled because D dimension is rounded up to 256 and runs inefficiently ++ //case GGML_TYPE_Q2_K: ++ //case GGML_TYPE_Q3_K: ++ //case GGML_TYPE_Q4_K: ++ //case GGML_TYPE_Q5_K: ++ //case GGML_TYPE_Q6_K: ++ case GGML_TYPE_IQ4_NL: ++ break; ++ default: ++ return false; ++ } ++ return true; ++ } ++ case GGML_OP_GET_ROWS: ++ { ++ switch (op->src[0]->type) { ++ case GGML_TYPE_F32: ++ case GGML_TYPE_F16: ++ case GGML_TYPE_Q4_0: ++ case GGML_TYPE_Q4_1: ++ case GGML_TYPE_Q5_0: ++ case GGML_TYPE_Q5_1: ++ case GGML_TYPE_Q8_0: ++ case GGML_TYPE_IQ4_NL: ++ return true; ++ default: ++ return false; ++ } ++ } break; ++ case GGML_OP_CONT: ++ case GGML_OP_CPY: ++ case GGML_OP_DUP: ++ { ++ ggml_type src0_type = op->src[0]->type; ++ ggml_type src1_type = op->src[1] != nullptr ? op->src[1]->type : src0_type; ++ if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { ++ return true; ++ } ++ if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F16) { ++ return true; ++ } ++ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { ++ return true; ++ } ++ return false; ++ } break; ++ case GGML_OP_REPEAT: ++ return ggml_type_size(op->type) == sizeof(float) && ggml_type_size(op->src[0]->type) == sizeof(float); ++ case GGML_OP_ROPE: ++ { ++ const int mode = ((const int32_t *) op->op_params)[2]; ++ if (mode & GGML_ROPE_TYPE_MROPE) { ++ return false; ++ } ++ if (mode & GGML_ROPE_TYPE_VISION) { ++ return false; ++ } ++ return ggml_is_contiguous(op->src[0]); ++ } ++ case GGML_OP_NONE: ++ case GGML_OP_RESHAPE: ++ case GGML_OP_VIEW: ++ case GGML_OP_PERMUTE: ++ case GGML_OP_TRANSPOSE: ++ case GGML_OP_NORM: ++ case GGML_OP_GROUP_NORM: ++ case GGML_OP_RMS_NORM: ++ case GGML_OP_ADD: ++ case GGML_OP_ACC: ++ case GGML_OP_MUL: ++ case GGML_OP_DIV: ++ case GGML_OP_CONCAT: ++ case GGML_OP_UPSCALE: ++ case GGML_OP_SCALE: ++ case GGML_OP_SQR: ++ case GGML_OP_SIN: ++ case GGML_OP_COS: ++ case GGML_OP_CLAMP: ++ case GGML_OP_PAD: ++ case GGML_OP_DIAG_MASK_INF: ++ case GGML_OP_SOFT_MAX: ++ case GGML_OP_ARGSORT: ++ case GGML_OP_SUM_ROWS: ++ case GGML_OP_IM2COL: ++ case GGML_OP_TIMESTEP_EMBEDDING: ++ case GGML_OP_POOL_2D: ++ case GGML_OP_RWKV_WKV6: ++ case GGML_OP_LEAKY_RELU: ++ return true; ++ default: ++ return false; ++ } ++ ++ UNUSED(dev); ++} ++ ++static bool ggml_backend_vk_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { ++ if (buft->iface.get_name != ggml_backend_vk_buffer_type_name) { ++ return false; ++ } ++ ++ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; ++ ggml_backend_vk_buffer_type_context * buft_ctx = (ggml_backend_vk_buffer_type_context *)buft->context; ++ ++ return buft_ctx->device->idx == ctx->device; ++} ++ ++static bool ggml_backend_vk_device_offload_op(ggml_backend_dev_t dev, const ggml_tensor * op) { ++ const int min_batch_size = 32; ++ ++ return (op->ne[1] >= min_batch_size && op->op != GGML_OP_GET_ROWS) || ++ (op->ne[2] >= min_batch_size && op->op == GGML_OP_MUL_MAT_ID); ++ ++ UNUSED(dev); ++} ++ ++static const struct ggml_backend_device_i ggml_backend_vk_device_i = { ++ /* .get_name = */ ggml_backend_vk_device_get_name, ++ /* .get_description = */ ggml_backend_vk_device_get_description, ++ /* .get_memory = */ ggml_backend_vk_device_get_memory, ++ /* .get_type = */ ggml_backend_vk_device_get_type, ++ /* .get_props = */ ggml_backend_vk_device_get_props, ++ /* .init_backend = */ ggml_backend_vk_device_init, ++ /* .get_buffer_type = */ ggml_backend_vk_device_get_buffer_type, ++ /* .get_host_buffer_type = */ ggml_backend_vk_device_get_host_buffer_type, ++ /* .buffer_from_host_ptr = */ NULL, ++ /* .supports_op = */ ggml_backend_vk_device_supports_op, ++ /* .supports_buft = */ ggml_backend_vk_device_supports_buft, ++ /* .offload_op = */ ggml_backend_vk_device_offload_op, ++ /* .event_new = */ NULL, ++ /* .event_free = */ NULL, ++ /* .event_synchronize = */ NULL, ++}; ++ ++static const char * ggml_backend_vk_reg_get_name(ggml_backend_reg_t reg) { ++ UNUSED(reg); ++ return GGML_VK_NAME; ++} ++ ++static size_t ggml_backend_vk_reg_get_device_count(ggml_backend_reg_t reg) { ++ UNUSED(reg); ++ return ggml_backend_vk_get_device_count(); ++} ++ ++static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg, size_t device) { ++ static std::vector devices; ++ ++ static bool initialized = false; ++ ++ { ++ static std::mutex mutex; ++ std::lock_guard lock(mutex); ++ if (!initialized) { ++ for (int i = 0; i < ggml_backend_vk_get_device_count(); i++) { ++ ggml_backend_vk_device_context * ctx = new ggml_backend_vk_device_context; ++ char desc[256]; ++ ggml_backend_vk_get_device_description(i, desc, sizeof(desc)); ++ ctx->device = i; ++ ctx->name = GGML_VK_NAME + std::to_string(i); ++ ctx->description = desc; ++ devices.push_back(new ggml_backend_device { ++ /* .iface = */ ggml_backend_vk_device_i, ++ /* .reg = */ reg, ++ /* .context = */ ctx, ++ }); ++ } ++ initialized = true; ++ } ++ } ++ ++ GGML_ASSERT(device < devices.size()); ++ return devices[device]; ++} ++ ++static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = { ++ /* .get_name = */ ggml_backend_vk_reg_get_name, ++ /* .get_device_count = */ ggml_backend_vk_reg_get_device_count, ++ /* .get_device = */ ggml_backend_vk_reg_get_device, ++ /* .get_proc_address = */ NULL, ++}; ++ ++ggml_backend_reg_t ggml_backend_vk_reg() { ++ static ggml_backend_reg reg = { ++ /* .api_version = */ GGML_BACKEND_API_VERSION, ++ /* .iface = */ ggml_backend_vk_reg_i, ++ /* .context = */ nullptr, ++ }; ++ ++ return ® ++} ++ ++// Extension availability ++static bool ggml_vk_instance_validation_ext_available(const std::vector& instance_extensions) { ++#ifdef GGML_VULKAN_VALIDATE ++ bool portability_enumeration_ext = false; ++ // Check for portability enumeration extension for MoltenVK support ++ for (const auto& properties : instance_extensions) { ++ if (strcmp("VK_KHR_portability_enumeration", properties.extensionName) == 0) { ++ return true; ++ } ++ } ++ if (!portability_enumeration_ext) { ++ std::cerr << "ggml_vulkan: WARNING: Instance extension VK_KHR_portability_enumeration not found." << std::endl; ++ } ++#endif ++ return false; ++ ++ UNUSED(instance_extensions); ++} ++static bool ggml_vk_instance_portability_enumeration_ext_available(const std::vector& instance_extensions) { ++#ifdef __APPLE__ ++ bool portability_enumeration_ext = false; ++ // Check for portability enumeration extension for MoltenVK support ++ for (const auto& properties : instance_extensions) { ++ if (strcmp("VK_KHR_portability_enumeration", properties.extensionName) == 0) { ++ return true; ++ } ++ } ++ if (!portability_enumeration_ext) { ++ std::cerr << "ggml_vulkan: WARNING: Instance extension VK_KHR_portability_enumeration not found." << std::endl; ++ } ++#endif ++ return false; ++ ++ UNUSED(instance_extensions); ++} ++ ++static bool ggml_vk_khr_cooperative_matrix_support(const vk::PhysicalDeviceProperties& props, const vk::PhysicalDeviceDriverProperties& driver_props) { ++ switch (props.vendorID) { ++ case VK_VENDOR_ID_INTEL: ++ // Intel drivers don't support coopmat properly yet ++ return false; ++ case VK_VENDOR_ID_AMD: ++ if (driver_props.driverID == vk::DriverId::eAmdProprietary || driver_props.driverID == vk::DriverId::eAmdOpenSource) { ++ // Workaround for AMD proprietary driver reporting support on all GPUs ++ const std::string name = props.deviceName; ++ return name.rfind("AMD Radeon RX 7", 0) == 0 || name.rfind("AMD Radeon(TM) RX 7", 0) == 0 || // RDNA 3 consumer GPUs ++ name.rfind("AMD Radeon PRO W7", 0) == 0 || name.rfind("AMD Radeon(TM) PRO W7", 0) == 0 || // RDNA 3 workstation GPUs ++ name.rfind("AMD Radeon 7", 0) == 0 || name.rfind("AMD Radeon(TM) 7", 0) == 0; // RDNA 3 APUs ++ } ++ return true; ++ default: ++ return true; ++ } ++} ++ ++// checks ++ ++#ifdef GGML_VULKAN_CHECK_RESULTS ++static void ggml_vk_print_graph_origin(const ggml_tensor * tensor, std::vector& done, int level = 0) { ++ if (std::find(done.begin(), done.end(), tensor) != done.end() || level > 10) { ++ return; ++ } ++ for (int j = 0; j < level; j++) { ++ std::cerr << " "; ++ } ++ std::cerr << ggml_op_name(tensor->op) << " gpu=" << (tensor->extra != nullptr) << std::endl; ++ ++ done.push_back(tensor); ++ ++ for (int i = 0; i < GGML_MAX_SRC; i++) { ++ if (tensor->src[i] != nullptr) { ++ ggml_vk_print_graph_origin(tensor->src[i], done, level + 1); ++ } ++ } ++} ++ ++static void ggml_vk_print_tensor_area(const ggml_tensor * tensor, const void * data, int i0, int i1, int i2, int i3) { ++ if (tensor->type != GGML_TYPE_F32 && tensor->type != GGML_TYPE_F16 && tensor->type != GGML_TYPE_I32) { ++ return; ++ } ++ i0 = std::max(i0, 5); ++ i1 = std::max(i1, 5); ++ i2 = std::max(i2, 0); ++ i3 = std::max(i3, 0); ++ fprintf(stderr, " "); ++ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { ++ fprintf(stderr, "%7d ", idx1); ++ } ++ fprintf(stderr, "\n"); ++ for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { ++ fprintf(stderr, "%7d: ", idx0); ++ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { ++ if (idx0 >= 0 && idx0 < tensor->ne[0] && idx1 >= 0 && idx1 < tensor->ne[1] && i2 >= 0 && i2 < tensor->ne[2] && i3 >= 0 && i3 < tensor->ne[3]) { ++ float val; ++ if (tensor->type == GGML_TYPE_F32) { ++ val = *(const float *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); ++ } else if (tensor->type == GGML_TYPE_F16) { ++ val = ggml_fp16_to_fp32(*(const ggml_fp16_t *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0])); ++ } else if (tensor->type == GGML_TYPE_I32) { ++ val = *(const int32_t *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ fprintf(stderr, "% 7.2f ", val); ++ } else { ++ fprintf(stderr, " "); ++ } ++ } ++ fprintf(stderr, "\n"); ++ } ++} ++ ++static void ggml_vk_print_tensor(const ggml_tensor * tensor, const char * name) { ++ void * tensor_data = tensor->data; ++ ++ const bool is_gpu = tensor->buffer != nullptr && ggml_backend_buffer_is_vk(tensor->buffer); ++ ++ if (is_gpu) { ++ const size_t tensor_size = ggml_nbytes(tensor); ++ tensor_data = malloc(tensor_size); ++ ++ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; ++ ++ vk_buffer buffer_gpu = buf_ctx->dev_buffer; ++ ggml_vk_buffer_read(buffer_gpu, vk_tensor_offset(tensor) + tensor->view_offs, tensor_data, tensor_size); ++ } ++ ++ std::cerr << "TENSOR CHECK " << name << " (" << tensor->name << "): " << ggml_op_name(tensor->op) << std::endl; ++ std::cerr << "tensor=" << tensor << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << std::endl; ++ if (tensor->src[0] != nullptr) { ++ std::cerr << "tensor->src[0]=" << tensor->src[0] << " name=" << tensor->src[0]->name << " op=" << ggml_op_name(tensor->src[0]->op) << " type=" << ggml_type_name(tensor->src[0]->type) << " ne0=" << tensor->src[0]->ne[0] << " nb0=" << tensor->src[0]->nb[0] << " ne1=" << tensor->src[0]->ne[1] << " nb1=" << tensor->src[0]->nb[1] << " ne2=" << tensor->src[0]->ne[2] << " nb2=" << tensor->src[0]->nb[2] << " ne3=" << tensor->src[0]->ne[3] << " nb3=" << tensor->src[0]->nb[3] << std::endl; ++ } ++ if (tensor->src[1] != nullptr) { ++ std::cerr << "tensor->src[1]=" << tensor->src[1] << " name=" << tensor->src[1]->name << " op=" << ggml_op_name(tensor->src[1]->op) << " type=" << ggml_type_name(tensor->src[1]->type) << " ne0=" << tensor->src[1]->ne[0] << " nb0=" << tensor->src[1]->nb[0] << " ne1=" << tensor->src[1]->ne[1] << " nb1=" << tensor->src[1]->nb[1] << " ne2=" << tensor->src[1]->ne[2] << " nb2=" << tensor->src[1]->nb[2] << " ne3=" << tensor->src[1]->ne[3] << " nb3=" << tensor->src[1]->nb[3] << std::endl; ++ } ++ std::cerr << std::endl << "Result:" << std::endl; ++ ggml_vk_print_tensor_area(tensor, tensor_data, 5, 5, 0, 0); ++ std::cerr << std::endl; ++ std::vector done; ++ ggml_vk_print_graph_origin(tensor, done); ++ ++ if (is_gpu) { ++ free(tensor_data); ++ } ++} ++ ++void * comp_result; ++size_t comp_size; ++size_t comp_nb[GGML_MAX_DIMS]; ++size_t check_counter = 0; ++static void ggml_vk_check_results_0(ggml_tensor * tensor) { ++ if (tensor->op == GGML_OP_TRANSPOSE) { ++ return; ++ } ++ ++ check_counter++; ++ if (!(vk_output_tensor > 0 && vk_output_tensor == check_counter) && check_counter <= vk_skip_checks) { ++ return; ++ } ++ ++ VK_LOG_DEBUG("ggml_vk_check_results_0(" << tensor->name << ")"); ++ ++ ggml_tensor * src0 = tensor->src[0]; ++ ggml_tensor * src1 = tensor->src[1]; ++ ggml_tensor * src2 = tensor->src[2]; ++ ggml_tensor * src3 = tensor->src[3]; ++ ++ struct ggml_init_params iparams = { ++ /*.mem_size =*/ 2ul*1024ul*1024ul*1024ul, ++ /*.mem_buffer =*/ NULL, ++ /*.no_alloc =*/ false, ++ }; ++ ++ struct ggml_context * ggml_ctx = ggml_init(iparams); ++ ++ struct ggml_tensor * src0_clone = nullptr; ++ struct ggml_tensor * src1_clone = nullptr; ++ struct ggml_tensor * src2_clone = nullptr; ++ struct ggml_tensor * src3_clone = nullptr; ++ struct ggml_tensor * tensor_clone = nullptr; ++ ++ size_t src0_size; ++ size_t src1_size; ++ size_t src2_size; ++ size_t src3_size; ++ ++ void * src0_buffer = nullptr; ++ void * src1_buffer = nullptr; ++ void * src2_buffer = nullptr; ++ void * src3_buffer = nullptr; ++ ++ if (src0 != nullptr) { ++ src0_clone = ggml_dup_tensor(ggml_ctx, src0); ++ ++ src0_size = ggml_nbytes(src0); ++ ++ src0_buffer = malloc(src0_size); ++ src0_clone->data = src0_buffer; ++ if (ggml_backend_buffer_is_host(src0->buffer)) { ++ memcpy(src0_clone->data, src0->data, src0_size); ++ memcpy(src0_clone->nb, src0->nb, sizeof(size_t) * GGML_MAX_DIMS); ++ } else if (ggml_backend_buffer_is_vk(src0->buffer)) { ++ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; ++ vk_buffer& buffer_gpu = buf_ctx->dev_buffer; ++ uint64_t offset = vk_tensor_offset(src0) + src0->view_offs; ++ if (!ggml_is_contiguous(src0) && ggml_vk_dim01_contiguous(src0)) { ++ for (int i3 = 0; i3 < src0->ne[3]; i3++) { ++ for (int i2 = 0; i2 < src0->ne[2]; i2++) { ++ const int idx = i3*src0->ne[2] + i2; ++ ggml_vk_buffer_read(buffer_gpu, offset + idx * src0->nb[2], ((char *)src0_clone->data + idx * src0_clone->nb[2]), src0->ne[1] * src0->nb[1]); ++ } ++ } ++ ++ src0_clone->nb[0] = src0->nb[0]; ++ src0_clone->nb[1] = src0->nb[1]; ++ for (int i = 2; i < GGML_MAX_DIMS; i++) { ++ src0_clone->nb[i] = src0_clone->nb[i - 1]*src0_clone->ne[i - 1]; ++ } ++ } else { ++ if (offset + src0_size >= buffer_gpu->size) { ++ src0_size = buffer_gpu->size - offset; ++ } ++ ggml_vk_buffer_read(buffer_gpu, offset, src0_clone->data, src0_size); ++ memcpy(src0_clone->nb, src0->nb, sizeof(size_t) * GGML_MAX_DIMS); ++ } ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ ++ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { ++ ggml_vk_print_tensor(src0, "src0"); ++ } ++ } ++ if (src1 != nullptr) { ++ src1_clone = ggml_dup_tensor(ggml_ctx, src1); ++ ++ src1_size = ggml_nbytes(src1); ++ ++ src1_buffer = malloc(src1_size); ++ src1_clone->data = src1_buffer; ++ if (ggml_backend_buffer_is_host(src1->buffer)) { ++ memcpy(src1_clone->data, src1->data, src1_size); ++ memcpy(src1_clone->nb, src1->nb, sizeof(size_t) * GGML_MAX_DIMS); ++ } else if (ggml_backend_buffer_is_vk(src1->buffer)) { ++ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; ++ vk_buffer& buffer_gpu = buf_ctx->dev_buffer; ++ uint64_t offset = vk_tensor_offset(src1) + src1->view_offs; ++ if (!ggml_is_contiguous(src1) && ggml_vk_dim01_contiguous(src1)) { ++ for (int i3 = 0; i3 < src1->ne[3]; i3++) { ++ for (int i2 = 0; i2 < src1->ne[2]; i2++) { ++ const int idx = i3*src1->ne[2] + i2; ++ ggml_vk_buffer_read(buffer_gpu, offset + idx * src1->nb[2], ((char *)src1_clone->data + idx * src1_clone->nb[2]), src1->ne[1] * src1->nb[1]); ++ } ++ } ++ ++ src1_clone->nb[0] = src1->nb[0]; ++ src1_clone->nb[1] = src1->nb[1]; ++ for (int i = 2; i < GGML_MAX_DIMS; i++) { ++ src1_clone->nb[i] = src1_clone->nb[i - 1]*src1_clone->ne[i - 1]; ++ } ++ } else { ++ if (offset + src1_size >= buffer_gpu->size) { ++ src1_size = buffer_gpu->size - offset; ++ } ++ ggml_vk_buffer_read(buffer_gpu, offset, src1_clone->data, src1_size); ++ memcpy(src1_clone->nb, src1->nb, sizeof(size_t) * GGML_MAX_DIMS); ++ } ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ ++ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { ++ ggml_vk_print_tensor(src1, "src1"); ++ } ++ } ++ if (src2 != nullptr) { ++ src2_clone = ggml_dup_tensor(ggml_ctx, src2); ++ ++ src2_size = ggml_nbytes(src2); ++ ++ src2_buffer = malloc(src2_size); ++ src2_clone->data = src2_buffer; ++ if (ggml_backend_buffer_is_host(src2->buffer)) { ++ memcpy(src2_clone->data, src2->data, src2_size); ++ memcpy(src2_clone->nb, src2->nb, sizeof(size_t) * GGML_MAX_DIMS); ++ } else if (ggml_backend_buffer_is_vk(src2->buffer)) { ++ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src2->buffer->context; ++ vk_buffer& buffer_gpu = buf_ctx->dev_buffer; ++ uint64_t offset = vk_tensor_offset(src2) + src2->view_offs; ++ if (!ggml_is_contiguous(src2) && ggml_vk_dim01_contiguous(src2)) { ++ for (int i3 = 0; i3 < src2->ne[3]; i3++) { ++ for (int i2 = 0; i2 < src2->ne[2]; i2++) { ++ const int idx = i3*src2->ne[2] + i2; ++ ggml_vk_buffer_read(buffer_gpu, offset + idx * src2->nb[2], ((char *)src2_clone->data + idx * src2_clone->nb[2]), src2->ne[1] * src2->nb[1]); ++ } ++ } ++ ++ src2_clone->nb[0] = src2->nb[0]; ++ src2_clone->nb[1] = src2->nb[1]; ++ for (int i = 2; i < GGML_MAX_DIMS; i++) { ++ src2_clone->nb[i] = src2_clone->nb[i - 1]*src2_clone->ne[i - 1]; ++ } ++ } else { ++ if (offset + src2_size >= buffer_gpu->size) { ++ src2_size = buffer_gpu->size - offset; ++ } ++ ggml_vk_buffer_read(buffer_gpu, offset, src2_clone->data, src2_size); ++ memcpy(src2_clone->nb, src2->nb, sizeof(size_t) * GGML_MAX_DIMS); ++ } ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ ++ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { ++ ggml_vk_print_tensor(src2, "src2"); ++ } ++ } ++ if (src3 != nullptr) { ++ src3_clone = ggml_dup_tensor(ggml_ctx, src3); ++ ++ src3_size = ggml_nbytes(src3); ++ ++ src3_buffer = malloc(src3_size); ++ src3_clone->data = src3_buffer; ++ if (ggml_backend_buffer_is_host(src3->buffer)) { ++ memcpy(src3_clone->data, src3->data, src3_size); ++ memcpy(src3_clone->nb, src3->nb, sizeof(size_t) * GGML_MAX_DIMS); ++ } else if (ggml_backend_buffer_is_vk(src3->buffer)) { ++ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src3->buffer->context; ++ vk_buffer& buffer_gpu = buf_ctx->dev_buffer; ++ uint64_t offset = vk_tensor_offset(src3) + src3->view_offs; ++ if (!ggml_is_contiguous(src3) && ggml_vk_dim01_contiguous(src3)) { ++ for (int i3 = 0; i3 < src3->ne[3]; i3++) { ++ for (int i2 = 0; i2 < src3->ne[2]; i2++) { ++ const int idx = i3*src3->ne[2] + i2; ++ ggml_vk_buffer_read(buffer_gpu, offset + idx * src3->nb[2], ((char *)src3_clone->data + idx * src3_clone->nb[2]), src3->ne[1] * src3->nb[1]); ++ } ++ } ++ ++ src3_clone->nb[0] = src3->nb[0]; ++ src3_clone->nb[1] = src3->nb[1]; ++ for (int i = 2; i < GGML_MAX_DIMS; i++) { ++ src3_clone->nb[i] = src3_clone->nb[i - 1]*src3_clone->ne[i - 1]; ++ } ++ } else { ++ if (offset + src3_size >= buffer_gpu->size) { ++ src3_size = buffer_gpu->size - offset; ++ } ++ ggml_vk_buffer_read(buffer_gpu, offset, src3_clone->data, src3_size); ++ memcpy(src3_clone->nb, src3->nb, sizeof(size_t) * GGML_MAX_DIMS); ++ } ++ } else { ++ GGML_ABORT("fatal error"); ++ } ++ ++ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { ++ ggml_vk_print_tensor(src3, "src3"); ++ } ++ } ++ ++ if (tensor->op == GGML_OP_FLASH_ATTN_EXT) { ++ const float *params = (const float *)tensor->op_params; ++ tensor_clone = ggml_flash_attn_ext(ggml_ctx, src0_clone, src1_clone, src2_clone, src3_clone, params[0], params[1], params[2]); ++ } else if (tensor->op == GGML_OP_MUL_MAT) { ++ tensor_clone = ggml_mul_mat(ggml_ctx, src0_clone, src1_clone); ++ } else if (tensor->op == GGML_OP_MUL_MAT_ID) { ++ tensor_clone = ggml_mul_mat_id(ggml_ctx, src0_clone, src1_clone, src2_clone); ++ } else if (tensor->op == GGML_OP_MUL) { ++ tensor_clone = ggml_mul(ggml_ctx, src0_clone, src1_clone); ++ } else if (tensor->op == GGML_OP_DIV) { ++ tensor_clone = ggml_div(ggml_ctx, src0_clone, src1_clone); ++ } else if (tensor->op == GGML_OP_CONCAT) { ++ tensor_clone = ggml_concat(ggml_ctx, src0_clone, src1_clone, *(int *)tensor->op_params); ++ } else if (tensor->op == GGML_OP_UPSCALE) { ++ tensor_clone = ggml_upscale_ext(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); ++ } else if (tensor->op == GGML_OP_SCALE) { ++ tensor_clone = ggml_scale(ggml_ctx, src0_clone, ((float *)tensor->op_params)[0]); ++ } else if (tensor->op == GGML_OP_SQR) { ++ tensor_clone = ggml_sqr(ggml_ctx, src0_clone); ++ } else if (tensor->op == GGML_OP_SIN) { ++ tensor_clone = ggml_sin(ggml_ctx, src0_clone); ++ } else if (tensor->op == GGML_OP_COS) { ++ tensor_clone = ggml_cos(ggml_ctx, src0_clone); ++ } else if (tensor->op == GGML_OP_CLAMP) { ++ tensor_clone = ggml_clamp(ggml_ctx, src0_clone, ((float *)tensor->op_params)[0], ((float *)tensor->op_params)[1]); ++ } else if (tensor->op == GGML_OP_PAD) { ++ tensor_clone = ggml_pad(ggml_ctx, src0_clone, tensor->ne[0] - src0_clone->ne[0], tensor->ne[1] - src0_clone->ne[1], tensor->ne[2] - src0_clone->ne[2], tensor->ne[3] - src0_clone->ne[3]); ++ } else if (tensor->op == GGML_OP_REPEAT) { ++ tensor_clone = ggml_repeat(ggml_ctx, src0_clone, tensor); ++ } else if (tensor->op == GGML_OP_ADD) { ++ tensor_clone = ggml_add(ggml_ctx, src0_clone, src1_clone); ++ } else if (tensor->op == GGML_OP_ACC) { ++ tensor_clone = ggml_acc(ggml_ctx, src0_clone, src1_clone, tensor->op_params[0], tensor->op_params[1], tensor->op_params[2], tensor->op_params[3]); ++ } else if (tensor->op == GGML_OP_NORM) { ++ tensor_clone = ggml_norm(ggml_ctx, src0_clone, *(float *)tensor->op_params); ++ } else if (tensor->op == GGML_OP_GROUP_NORM) { ++ tensor_clone = ggml_group_norm(ggml_ctx, src0_clone, *(int *)tensor->op_params, ((float *)tensor->op_params)[1]); ++ } else if (tensor->op == GGML_OP_RMS_NORM) { ++ tensor_clone = ggml_rms_norm(ggml_ctx, src0_clone, *(float *)tensor->op_params); ++ } else if (tensor->op == GGML_OP_SOFT_MAX) { ++ if (src1 != nullptr) { ++ tensor_clone = ggml_soft_max_ext(ggml_ctx, src0_clone, src1_clone, ((float *)tensor->op_params)[0], ((float *)tensor->op_params)[1]); ++ } else { ++ tensor_clone = ggml_soft_max(ggml_ctx, src0_clone); ++ } ++ } else if (tensor->op == GGML_OP_DIAG_MASK_INF) { ++ tensor_clone = ggml_diag_mask_inf(ggml_ctx, src0_clone, *(int *)tensor->op_params); ++ } else if (tensor->op == GGML_OP_ROPE) { ++ const int n_dims = ((int32_t *) tensor->op_params)[1]; ++ const int mode = ((int32_t *) tensor->op_params)[2]; ++ //const int n_ctx_ggml = ((int32_t *) tensor->op_params)[3]; ++ const int n_ctx_orig_ggml = ((int32_t *) tensor->op_params)[4]; ++ const float freq_base = ((float *) tensor->op_params)[5]; ++ const float freq_scale = ((float *) tensor->op_params)[6]; ++ const float ext_factor = ((float *) tensor->op_params)[7]; ++ const float attn_factor = ((float *) tensor->op_params)[8]; ++ const float beta_fast = ((float *) tensor->op_params)[9]; ++ const float beta_slow = ((float *) tensor->op_params)[10]; ++ tensor_clone = ggml_rope_ext(ggml_ctx, src0_clone, src1_clone, src2_clone, n_dims, mode, n_ctx_orig_ggml, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); ++ } else if (tensor->op == GGML_OP_UNARY) { ++ switch (ggml_get_unary_op(tensor)) { ++ case GGML_UNARY_OP_SILU: ++ tensor_clone = ggml_silu(ggml_ctx, src0_clone); ++ break; ++ case GGML_UNARY_OP_GELU: ++ tensor_clone = ggml_gelu(ggml_ctx, src0_clone); ++ break; ++ case GGML_UNARY_OP_GELU_QUICK: ++ tensor_clone = ggml_gelu_quick(ggml_ctx, src0_clone); ++ break; ++ case GGML_UNARY_OP_RELU: ++ tensor_clone = ggml_relu(ggml_ctx, src0_clone); ++ break; ++ case GGML_UNARY_OP_TANH: ++ tensor_clone = ggml_tanh(ggml_ctx, src0_clone); ++ break; ++ default: ++ std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl; ++ GGML_ABORT("fatal error"); ++ } ++ } else if (tensor->op == GGML_OP_CPY || tensor->op == GGML_OP_DUP) { ++ if (src1 == nullptr) { ++ tensor_clone = ggml_dup(ggml_ctx, src0_clone); ++ tensor_clone->type = tensor->type; ++ } else { ++ tensor_clone = ggml_cpy(ggml_ctx, src0_clone, src1_clone); ++ } ++ } else if (tensor->op == GGML_OP_CONT) { ++ tensor_clone = ggml_cont_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); ++ } else if (tensor->op == GGML_OP_RESHAPE) { ++ tensor_clone = ggml_reshape_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); ++ } else if (tensor->op == GGML_OP_VIEW) { ++ tensor_clone = ggml_view_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3], tensor->nb[1], tensor->nb[2], tensor->nb[3], ((int32_t *) tensor->op_params)[0]); ++ } else if (tensor->op == GGML_OP_PERMUTE) { ++ int32_t * params = (int32_t *)tensor->op_params; ++ tensor_clone = ggml_permute(ggml_ctx, src0_clone, params[0], params[1], params[2], params[3]); ++ } else if (tensor->op == GGML_OP_TRANSPOSE) { ++ tensor_clone = ggml_transpose(ggml_ctx, src0_clone); ++ } else if (tensor->op == GGML_OP_GET_ROWS) { ++ tensor_clone = ggml_get_rows(ggml_ctx, src0_clone, src1_clone); ++ } else if (tensor->op == GGML_OP_ARGSORT) { ++ tensor_clone = ggml_argsort(ggml_ctx, src0_clone, (ggml_sort_order) *(int *)tensor->op_params); ++ } else if (tensor->op == GGML_OP_SUM_ROWS) { ++ tensor_clone = ggml_sum_rows(ggml_ctx, src0_clone); ++ } else if (tensor->op == GGML_OP_IM2COL) { ++ const int32_t s0 = tensor->op_params[0]; ++ const int32_t s1 = tensor->op_params[1]; ++ const int32_t p0 = tensor->op_params[2]; ++ const int32_t p1 = tensor->op_params[3]; ++ const int32_t d0 = tensor->op_params[4]; ++ const int32_t d1 = tensor->op_params[5]; ++ ++ const bool is_2D = tensor->op_params[6] == 1; ++ tensor_clone = ggml_im2col(ggml_ctx, src0_clone, src1_clone, s0, s1, p0, p1, d0, d1, is_2D, tensor->type); ++ } else if (tensor->op == GGML_OP_TIMESTEP_EMBEDDING) { ++ const int32_t dim = tensor->op_params[0]; ++ const int32_t max_period = tensor->op_params[1]; ++ tensor_clone = ggml_timestep_embedding(ggml_ctx, src0_clone, dim, max_period); ++ } else if (tensor->op == GGML_OP_POOL_2D) { ++ enum ggml_op_pool op = static_cast(tensor->op_params[0]); ++ const int32_t k0 = tensor->op_params[1]; ++ const int32_t k1 = tensor->op_params[2]; ++ const int32_t s0 = tensor->op_params[3]; ++ const int32_t s1 = tensor->op_params[4]; ++ const int32_t p0 = tensor->op_params[5]; ++ const int32_t p1 = tensor->op_params[6]; ++ ++ tensor_clone = ggml_pool_2d(ggml_ctx, src0_clone, op, k0, k1, s0, s1, p0, p1); ++ } else if (tensor->op == GGML_OP_LEAKY_RELU) { ++ const float * op_params = (const float *)tensor->op_params; ++ tensor_clone = ggml_leaky_relu(ggml_ctx, src0_clone, op_params[0], false); ++ } else if (tensor->op == GGML_OP_RWKV_WKV6) { ++ tensor_clone = ggml_rwkv_wkv6(ggml_ctx, tensor->src[0], tensor->src[1], tensor->src[2], tensor->src[3], ++ tensor->src[4], tensor->src[5]); ++ } ++ else { ++ std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl; ++ GGML_ABORT("fatal error"); ++ } ++ ++ ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); ++ ggml_build_forward_expand(cgraph, tensor_clone); ++ ++ ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 8); ++ ++ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { ++ ggml_vk_print_tensor(tensor_clone, "tensor_clone"); ++ } ++ ++ comp_size = ggml_nbytes(tensor_clone); ++ ++ comp_result = malloc(comp_size); ++ memcpy(comp_result, tensor_clone->data, comp_size); ++ memcpy(comp_nb, tensor_clone->nb, sizeof(size_t) * GGML_MAX_DIMS); ++ ++ if (src0 != nullptr) { ++ free(src0_buffer); ++ } ++ if (src1 != nullptr) { ++ free(src1_buffer); ++ } ++ ++ ggml_free(ggml_ctx); ++ ++ VK_LOG_DEBUG("END ggml_vk_check_results_0(" << tensor->name << ")"); ++} ++ ++static void ggml_vk_check_results_1(ggml_tensor * tensor) { ++ if (tensor->op == GGML_OP_TRANSPOSE) { ++ return; ++ } ++ if (!(vk_output_tensor > 0 && vk_output_tensor == check_counter) && check_counter <= vk_skip_checks) { ++ return; ++ } ++ ++ VK_LOG_DEBUG("ggml_vk_check_results_1(" << tensor->name << ")"); ++ ++ ggml_tensor * src0 = tensor->src[0]; ++ ggml_tensor * src1 = tensor->src[1]; ++ ggml_tensor * src2 = tensor->src[2]; ++ ++ void * tensor_data = tensor->data; ++ ++ if (ggml_backend_buffer_is_vk(tensor->buffer)) { ++ size_t tensor_size = ggml_nbytes(tensor); ++ tensor_data = malloc(tensor_size); ++ ++ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; ++ ++ vk_buffer& buffer_gpu = buf_ctx->dev_buffer; ++ uint64_t offset = vk_tensor_offset(tensor) + tensor->view_offs; ++ if (offset + tensor_size >= buffer_gpu->size) { ++ tensor_size = buffer_gpu->size - offset; ++ } ++ ++ ggml_vk_buffer_read(buffer_gpu, offset, tensor_data, tensor_size); ++ } ++ ++ float first_error_result = -1.0f; ++ float first_error_correct = -1.0f; ++ std::array first_error = { -1, -1, -1, -1 }; ++ double avg_err = 0.0; ++ size_t counter = 0; ++ ++ for (int i3 = 0; i3 < tensor->ne[3]; i3++) { ++ for (int i2 = 0; i2 < tensor->ne[2]; i2++) { ++ for (int i1 = 0; i1 < tensor->ne[1]; i1++) { ++ for (int i0 = 0; i0 < tensor->ne[0]; i0++) { ++ const bool buffer_size_fit = i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0] < comp_size; ++ float correct = 0.0f; ++ float result = 0.0f; ++ ++ if (buffer_size_fit) { ++ if (tensor->type == GGML_TYPE_F32) { ++ correct = *(float *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0]); ++ result = *(float *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0]); ++ } else if (tensor->type == GGML_TYPE_F16) { ++ correct = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0])); ++ result = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0])); ++ } else if (tensor->type == GGML_TYPE_I32) { ++ correct = *(int32_t *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0]); ++ result = *(int32_t *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0]); ++ } else { ++ std::cerr << "Results check not implemented for type " << ggml_type_name(tensor->type) << std::endl; ++ } ++ } else { ++ std::cerr << "Missing debug code for type " << ggml_type_name(tensor->type) << std::endl; ++ GGML_ABORT("fatal error"); ++ } ++ ++ if ((std::isnan(correct) != std::isnan(result)) || (std::isinf(correct) != std::isinf(result)) || !buffer_size_fit) { ++ std::cerr << "ERROR: Invalid value in " << ggml_op_name(tensor->op) << " i3=" << i3 << " i2=" << i2 << " i1=" << i1 << " i0=" << i0 << " result=" << result << " correct=" << correct << " avg_err=" << (avg_err / counter) << std::endl; ++ std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; ++ if (src0 != nullptr) { ++ std::cerr << "src0=" << src0 << " src0->name=" << src0->name << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; ++ } ++ if (src1 != nullptr) { ++ std::cerr << "src1=" << src1 << " src1->name=" << src1->name << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; ++ } ++ if (src2 != nullptr) { ++ std::cerr << "src2=" << src2 << " src2->name=" << src2->name << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; ++ } ++ std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; ++ std::cerr << std::endl << "Result:" << std::endl; ++ ggml_vk_print_tensor_area(tensor, tensor_data, i0, i1, i2, i3); ++ std::cerr << std::endl << "Correct:" << std::endl; ++ ggml_vk_print_tensor_area(tensor, comp_result, i0, i1, i2, i3); ++ std::cerr << std::endl; ++ std::vector done; ++ ggml_vk_print_graph_origin(tensor, done); ++ GGML_ABORT("fatal error"); ++ } ++ if (first_error[0] == -1 && std::fabs(correct - result) > 0.1f) { ++ first_error[0] = i0; ++ first_error[1] = i1; ++ first_error[2] = i2; ++ first_error[3] = i3; ++ first_error_result = result; ++ first_error_correct = correct; ++ } ++ ++ // Special case, value is infinite, avoid NaN result in avg_err ++ // NaN also appears in results, if both are nan error is 0 ++ if (!std::isinf(correct) && !std::isinf(result) && !std::isnan(correct) && !std::isnan(result)) { ++ avg_err += std::fabs(correct - result); ++ } ++ counter++; ++ } ++ } ++ } ++ } ++ ++ avg_err /= counter; ++ ++ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { ++ std::cerr << "TENSOR CHECK: avg_err=" << avg_err << " in " << ggml_op_name(tensor->op) << " (check " << check_counter << ")" << std::endl; ++ std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; ++ if (src0 != nullptr) { ++ std::cerr << "src0=" << src0 << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; ++ } ++ if (src1 != nullptr) { ++ std::cerr << "src1=" << src1 << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; ++ } ++ if (src2 != nullptr) { ++ std::cerr << "src2=" << src2 << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; ++ } ++ std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; ++ std::cerr << std::endl << "Result:" << std::endl; ++ ggml_vk_print_tensor_area(tensor, tensor_data, 5, 5, 0, 0); ++ std::cerr << std::endl << "Correct:" << std::endl; ++ ggml_vk_print_tensor_area(tensor, comp_result, 5, 5, 0, 0); ++ std::cerr << std::endl; ++ std::vector done; ++ ggml_vk_print_graph_origin(tensor, done); ++ } ++ ++ if (avg_err > 0.05 || std::isnan(avg_err)) { ++ std::cerr << "ERROR: avg_err=" << avg_err << " in " << ggml_op_name(tensor->op) << " (check " << check_counter << ")" << std::endl; ++ std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; ++ if (src0 != nullptr) { ++ std::cerr << "src0=" << src0 << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; ++ } ++ if (src1 != nullptr) { ++ std::cerr << "src1=" << src1 << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; ++ } ++ if (src2 != nullptr) { ++ std::cerr << "src2=" << src2 << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; ++ } ++ std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; ++ std::cerr << std::endl << "Result:" << std::endl; ++ ggml_vk_print_tensor_area(tensor, tensor_data, first_error[0], first_error[1], first_error[2], first_error[3]); ++ std::cerr << std::endl << "Correct:" << std::endl; ++ ggml_vk_print_tensor_area(tensor, comp_result, first_error[0], first_error[1], first_error[2], first_error[3]); ++ std::cerr << std::endl; ++ std::vector done; ++ ggml_vk_print_graph_origin(tensor, done); ++ GGML_ABORT("fatal error"); ++ } else { ++ std::cerr << check_counter << " " << tensor->name << " op=" << ggml_op_name(tensor->op) << " avg_err=" << avg_err << std::endl; ++ } ++ ++ free(comp_result); ++ comp_result = nullptr; ++ comp_size = 0; ++ ++ if (ggml_backend_buffer_is_vk(tensor->buffer)) { ++ free(tensor_data); ++ } ++ ++ VK_LOG_DEBUG("END ggml_vk_check_results_1(" << tensor->name << ")"); ++} ++#endif ++ ++GGML_BACKEND_DL_IMPL(ggml_backend_vk_reg) +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt +new file mode 100644 +index 00000000..bd0c74cb +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt +@@ -0,0 +1,9 @@ ++find_package (Threads REQUIRED) ++find_package(Vulkan COMPONENTS glslc REQUIRED) ++ ++set(TARGET vulkan-shaders-gen) ++add_executable(${TARGET} vulkan-shaders-gen.cpp) ++install(TARGETS ${TARGET} RUNTIME) ++target_compile_features(${TARGET} PRIVATE cxx_std_17) ++target_link_libraries(vulkan-shaders-gen PUBLIC Threads::Threads) ++target_link_libraries(vulkan-shaders-gen PRIVATE Vulkan::Vulkan) +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp +new file mode 100644 +index 00000000..d896f1ef +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp +@@ -0,0 +1,29 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_binary_head.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ const uint idx = gl_GlobalInvocationID.x; ++ if (idx >= p.ne) { ++ return; ++ } ++ ++ const uint offset = p.param3; ++ const uint src1_i = idx - offset; ++ const uint oz = src1_i / p.nb02; ++ const uint oy = (src1_i - (oz * p.nb02)) / p.nb01; ++ const uint ox = src1_i % p.nb01; ++ ++ uint i00, i01, i02, i03; ++ get_indices(idx, i00, i01, i02, i03); ++ ++ if (ox < p.ne10 && oy < p.ne11 && oz < p.ne12) { ++ data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) + FLOAT_TYPE(data_b[get_boffset() + ox + oy * p.ne10 + oz * p.ne10 * p.ne11])); ++ } else { ++ data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)])); ++ } ++} ++ +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp +new file mode 100644 +index 00000000..2b4085c4 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp +@@ -0,0 +1,29 @@ ++#version 450 ++ ++#extension GL_EXT_shader_16bit_storage : require ++ ++#include "types.comp" ++#include "generic_binary_head.comp" ++ ++const uint num_threads = 256; ++ ++layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ uint idx = get_idx(); ++ ++ // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation ++ const uint num_iter = 2; ++ ++ [[unroll]] for (uint i = 0; i < num_iter; ++i) { ++ if (idx >= p.ne) { ++ continue; ++ } ++ uint i00, i01, i02, i03; ++ get_indices(idx, i00, i01, i02, i03); ++ ++ data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) + FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); ++ ++ idx += num_threads; ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp +new file mode 100644 +index 00000000..d4fa45b1 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp +@@ -0,0 +1,69 @@ ++#version 450 ++ ++#include "types.comp" ++ ++#define BLOCK_SIZE 1024 ++#define ASC 0 ++ ++layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) buffer D {int data_d[];}; ++ ++layout (push_constant) uniform parameter { ++ uint ncols; ++ uint ncols_pad; ++ uint order; ++} p; ++ ++shared int dst_row[BLOCK_SIZE]; ++ ++void swap(uint idx0, uint idx1) { ++ int tmp = dst_row[idx0]; ++ dst_row[idx0] = dst_row[idx1]; ++ dst_row[idx1] = tmp; ++} ++ ++void main() { ++ // bitonic sort ++ const int col = int(gl_LocalInvocationID.x); ++ const uint row = gl_WorkGroupID.y; ++ ++ const uint row_offset = row * p.ncols; ++ ++ // initialize indices ++ if (col < p.ncols_pad) { ++ dst_row[col] = col; ++ } ++ barrier(); ++ ++ for (uint k = 2; k <= p.ncols_pad; k *= 2) { ++ for (uint j = k / 2; j > 0; j /= 2) { ++ const uint ixj = col ^ j; ++ if (col < p.ncols_pad && ixj > col) { ++ if ((col & k) == 0) { ++ if (dst_row[col] >= p.ncols || ++ (dst_row[ixj] < p.ncols && (p.order == ASC ? ++ data_a[row_offset + dst_row[col]] > data_a[row_offset + dst_row[ixj]] : ++ data_a[row_offset + dst_row[col]] < data_a[row_offset + dst_row[ixj]])) ++ ) { ++ swap(col, ixj); ++ } ++ } else { ++ if (dst_row[ixj] >= p.ncols || ++ (dst_row[col] < p.ncols && (p.order == ASC ? ++ data_a[row_offset + dst_row[col]] < data_a[row_offset + dst_row[ixj]] : ++ data_a[row_offset + dst_row[col]] > data_a[row_offset + dst_row[ixj]])) ++ ) { ++ swap(col, ixj); ++ } ++ } ++ } ++ barrier(); ++ } ++ } ++ ++ if (col < p.ncols) { ++ data_d[row_offset + col] = dst_row[col]; ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp +new file mode 100644 +index 00000000..1e5cb8da +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp +@@ -0,0 +1,17 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_unary_head.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ const uint idx = get_idx(); ++ ++ if (idx >= p.ne) { ++ return; ++ } ++ ++ const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); ++ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(val < p.param1 ? p.param1 : (val > p.param2 ? p.param2 : val)); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp +new file mode 100644 +index 00000000..9ee2f1fa +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp +@@ -0,0 +1,41 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_binary_head.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; ++ const int dim = p.param3; ++ ++ if (idx >= p.ne) { ++ return; ++ } ++ ++ const uint i3 = idx / (p.ne22*p.ne21*p.ne20); ++ const uint i3_offset = i3 * p.ne22*p.ne21*p.ne20; ++ const uint i2 = (idx - i3_offset) / (p.ne21*p.ne20); ++ const uint i2_offset = i2*p.ne21*p.ne20; ++ const uint i1 = (idx - i3_offset - i2_offset) / p.ne20; ++ const uint i0 = idx - i3_offset - i2_offset - i1*p.ne20; ++ ++ uint o[4] = {0, 0, 0, 0}; ++ o[dim] = dim == 0 ? p.ne00 : (dim == 1 ? p.ne01 : (dim == 2 ? p.ne02 : p.ne03)); ++ ++ const uint src0_idx = i3*p.nb03 + i2*p.nb02 + i1*p.nb01 + i0*p.nb00; ++ const uint src1_idx = (i3 - o[3])*p.nb13 + (i2 - o[2])*p.nb12 + (i1 - o[1])*p.nb11 + (i0 - o[0])*p.nb10; ++ const uint dst_idx = i3*p.nb23 + i2*p.nb22 + i1*p.nb21 + i0*p.nb20; ++ ++ const bool is_src0 = i0 < p.ne00 && i1 < p.ne01 && i2 < p.ne02 && i3 < p.ne03; ++ ++#ifndef OPTIMIZATION_ERROR_WORKAROUND ++ data_d[get_doffset() + dst_idx] = D_TYPE(is_src0 ? data_a[get_aoffset() + src0_idx] : data_b[get_boffset() + src1_idx]); ++#else ++ if (is_src0) { ++ data_d[get_doffset() + dst_idx] = data_a[get_aoffset() + src0_idx]; ++ } else { ++ data_d[get_doffset() + dst_idx] = data_b[get_boffset() + src1_idx]; ++ } ++#endif ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp +new file mode 100644 +index 00000000..dd828c23 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp +@@ -0,0 +1,42 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_unary_head.comp" ++ ++#extension GL_EXT_control_flow_attributes : require ++ ++const uint num_threads = 128; ++ ++layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ uint idx = get_idx(); ++ ++ // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation ++ const uint num_iter = 4; ++ ++ // fast path for when all four iterations are in-bounds ++ if (idx + (num_iter-1)*num_threads < p.ne) { ++ [[unroll]] for (uint i = 0; i < num_iter; ++i) { ++#ifndef OPTIMIZATION_ERROR_WORKAROUND ++ data_d[get_doffset() + idx] = D_TYPE(data_a[get_aoffset() + idx]); ++#else ++ data_d[get_doffset() + idx] = data_a[get_aoffset() + idx]; ++#endif ++ idx += num_threads; ++ } ++ } else { ++ [[unroll]] for (uint i = 0; i < num_iter; ++i) { ++ if (idx >= p.ne) { ++ continue; ++ } ++ ++#ifndef OPTIMIZATION_ERROR_WORKAROUND ++ data_d[get_doffset() + idx] = D_TYPE(data_a[get_aoffset() + idx]); ++#else ++ data_d[get_doffset() + idx] = data_a[get_aoffset() + idx]; ++#endif ++ idx += num_threads; ++ } ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp +new file mode 100644 +index 00000000..29c90649 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp +@@ -0,0 +1,20 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_unary_head.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ const uint idx = get_idx(); ++ ++ if (idx >= p.ne) { ++ return; ++ } ++ ++#ifndef OPTIMIZATION_ERROR_WORKAROUND ++ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(data_a[get_aoffset() + src0_idx(idx)]); ++#else ++ data_d[get_doffset() + dst_idx(idx)] = data_a[get_aoffset() + src0_idx(idx)]; ++#endif ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp +new file mode 100644 +index 00000000..0b8d02f5 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp +@@ -0,0 +1,17 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_unary_head.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ const uint idx = get_idx(); ++ ++ if (idx >= p.ne) { ++ return; ++ } ++ ++ const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); ++ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(cos(val)); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp +new file mode 100644 +index 00000000..a4d3fca5 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp +@@ -0,0 +1,20 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {float data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ const uint i = gl_GlobalInvocationID.x * 16; ++ ++ if (i >= p.nel) { ++ return; ++ } ++ ++ [[unroll]] for (uint l = 0; l < 16; l++) { ++ data_b[i + l] = D_TYPE(data_a[i + l]); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp +new file mode 100644 +index 00000000..91bb8f8d +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp +@@ -0,0 +1,118 @@ ++#if !defined(DATA_A_F32) && !defined(DATA_A_F16) ++#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require ++#endif ++ ++#include "types.comp" ++ ++#if defined(A_TYPE_PACKED16) ++layout (binding = 0) readonly buffer A_PACKED16 {A_TYPE_PACKED16 data_a_packed16[];}; ++#endif ++#if defined(A_TYPE_PACKED32) ++layout (binding = 0) readonly buffer A_PACKED32 {A_TYPE_PACKED32 data_a_packed32[];}; ++#endif ++ ++#if defined(DATA_A_F32) ++vec2 dequantize(uint ib, uint iqs, uint a_offset) { ++ return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]); ++} ++#endif ++ ++#if defined(DATA_A_F16) ++vec2 dequantize(uint ib, uint iqs, uint a_offset) { ++ return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]); ++} ++#endif ++ ++#if defined(DATA_A_Q4_0) ++vec2 dequantize(uint ib, uint iqs, uint a_offset) { ++ const uint vui = uint(data_a[a_offset + ib].qs[iqs]); ++ return (vec2(vui & 0xF, vui >> 4) - 8.0f); ++} ++vec4 dequantize4(uint ib, uint iqs, uint a_offset) { ++ const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); ++ return (vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, vui >> 12) - 8.0f); ++} ++#endif ++ ++#if defined(DATA_A_Q4_1) ++vec2 dequantize(uint ib, uint iqs, uint a_offset) { ++ const uint vui = uint(data_a[a_offset + ib].qs[iqs]); ++ return vec2(vui & 0xF, vui >> 4); ++} ++vec4 dequantize4(uint ib, uint iqs, uint a_offset) { ++ const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); ++ return vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, vui >> 12); ++} ++#endif ++ ++#if defined(DATA_A_Q5_0) ++vec2 dequantize(uint ib, uint iqs, uint a_offset) { ++ const uint uint_qh = uint(data_a[a_offset + ib].qh[1]) << 16 | data_a[a_offset + ib].qh[0]; ++ const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); ++ const uint vui = uint(data_a[a_offset + ib].qs[iqs]); ++ return (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f); ++} ++vec4 dequantize4(uint ib, uint iqs, uint a_offset) { ++ const uint uint_qh = uint(data_a_packed16[a_offset + ib].qh[1]) << 16 | data_a_packed16[a_offset + ib].qh[0]; ++ const ivec2 qh0 = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); ++ const ivec2 qh1 = ivec2(((uint_qh >> (iqs + 1)) << 4) & 0x10, (uint_qh >> (iqs + 13)) & 0x10); ++ const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); ++ return (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f); ++} ++#endif ++ ++#if defined(DATA_A_Q5_1) ++vec2 dequantize(uint ib, uint iqs, uint a_offset) { ++ const uint uint_qh = data_a[a_offset + ib].qh; ++ const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); ++ const uint vui = uint(data_a[a_offset + ib].qs[iqs]); ++ return vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y); ++} ++vec4 dequantize4(uint ib, uint iqs, uint a_offset) { ++ const uint uint_qh = data_a_packed16[a_offset + ib].qh; ++ const ivec2 qh0 = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); ++ const ivec2 qh1 = ivec2(((uint_qh >> (iqs + 1)) << 4) & 0x10, (uint_qh >> (iqs + 13)) & 0x10); ++ const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); ++ return vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y); ++} ++#endif ++ ++#if defined(DATA_A_Q8_0) ++vec2 dequantize(uint ib, uint iqs, uint a_offset) { ++ return vec2(int(data_a[a_offset + ib].qs[iqs]), int(data_a[a_offset + ib].qs[iqs + 1])); ++} ++vec4 dequantize4(uint ib, uint iqs, uint a_offset) { ++ uint32_t v0 = data_a_packed16[a_offset + ib].qs[iqs/2]; ++ uint32_t v1 = data_a_packed16[a_offset + ib].qs[iqs/2 + 1]; ++ return vec4(int8_t(v0 & 0xFF), int8_t(v0 >> 8), int8_t(v1 & 0xFF), int8_t(v1 >> 8)); ++} ++#endif ++ ++#if defined(DATA_A_IQ4_NL) ++vec2 dequantize(uint ib, uint iqs, uint a_offset) { ++ const uint vui = uint(data_a[a_offset + ib].qs[iqs]); ++ return vec2(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[vui >> 4]); ++} ++vec4 dequantize4(uint ib, uint iqs, uint a_offset) { ++ const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); ++ return vec4(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[(vui >> 4) & 0xF], kvalues_iq4nl[(vui >> 8) & 0xF], kvalues_iq4nl[vui >> 12]); ++} ++#endif ++ ++#if defined(DATA_A_F32) || defined(DATA_A_F16) ++vec2 get_dm(uint ib, uint a_offset) { ++ return vec2(0, 0); ++} ++#endif ++ ++#if defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ4_NL) ++vec2 get_dm(uint ib, uint a_offset) { ++ return vec2(float(data_a[a_offset + ib].d), 0); ++} ++#endif ++ ++#if defined(DATA_A_Q4_1) || defined(DATA_A_Q5_1) ++vec2 get_dm(uint ib, uint a_offset) { ++ return vec2(float(data_a[a_offset + ib].d), float(data_a[a_offset + ib].m)); ++} ++#endif +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp +new file mode 100644 +index 00000000..94b78598 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp +@@ -0,0 +1,325 @@ ++ ++#include "types.comp" ++ ++layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ4_0 { ++ block_q4_0_packed16 block; ++}; ++ ++float16_t dequantFuncQ4_0(const in decodeBufQ4_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ const float16_t d = bl.block.d; ++ const uint idx = coordInBlock[1]; ++ const uint shift = (idx & 0x10) >> 2; ++ uint32_t qs = uint32_t(bl.block.qs[(idx & 0xE) >> 1]); ++ qs >>= shift; ++ qs &= 0x0F0F; ++ qs = unpack8(qs)[idx & 1]; ++ float16_t ret = (float16_t(qs) - float16_t(8)) * d; ++ return ret; ++} ++ ++layout(buffer_reference, std430, buffer_reference_align = 4) buffer decodeBufQ4_1 { ++ block_q4_1 block; ++}; ++ ++float16_t dequantFuncQ4_1(const in decodeBufQ4_1 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ const float16_t d = bl.block.d; ++ const float16_t m = bl.block.m; ++ const uint idx = coordInBlock[1]; ++ const uint iqs = idx & 0xF; ++ const uint shift = (idx & 0x10) >> 2; ++ uint32_t qs = bl.block.qs[iqs]; ++ qs >>= shift; ++ qs &= 0xF; ++ float16_t ret = float16_t(qs) * d + m; ++ return ret; ++} ++ ++layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ5_0 { ++ block_q5_0 block; ++}; ++ ++float16_t dequantFuncQ5_0(const in decodeBufQ5_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ const float16_t d = bl.block.d; ++ const uint idx = coordInBlock[1]; ++ const uint iqs = idx & 0xF; ++ ++ const uint uint_qh = uint(bl.block.qh[1]) << 16 | bl.block.qh[0]; ++ const uint qh = ((uint_qh >> idx) << 4) & 0x10; ++ ++ const uint shift = (idx & 0x10) >> 2; ++ uint32_t qs = bl.block.qs[iqs]; ++ qs >>= shift; ++ qs &= 0xF; ++ ++ float16_t ret = (float16_t(qs | qh) - float16_t(16)) * d; ++ return ret; ++} ++ ++layout(buffer_reference, std430, buffer_reference_align = 8) buffer decodeBufQ5_1 { ++ block_q5_1 block; ++}; ++ ++float16_t dequantFuncQ5_1(const in decodeBufQ5_1 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ const float16_t d = bl.block.d; ++ const float16_t m = bl.block.m; ++ const uint idx = coordInBlock[1]; ++ const uint iqs = idx & 0xF; ++ ++ const uint uint_qh = bl.block.qh; ++ const uint qh = ((uint_qh >> idx) << 4) & 0x10; ++ ++ const uint shift = (idx & 0x10) >> 2; ++ uint32_t qs = bl.block.qs[iqs]; ++ qs >>= shift; ++ qs &= 0xF; ++ ++ float16_t ret = float16_t(qs | qh) * d + m; ++ return ret; ++} ++ ++layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ8_0 { ++ block_q8_0_packed16 block; ++}; ++ ++float16_t dequantFuncQ8_0(const in decodeBufQ8_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ const float16_t d = bl.block.d; ++ const uint idx = coordInBlock[1]; ++ const uint iqs = idx; ++ ++ // Load 16b and select the byte for this element ++ int32_t qs = unpack8(int32_t(bl.block.qs[(iqs & 0x1E) >> 1]))[iqs & 1]; ++ float16_t ret = float16_t(qs) * d; ++ return ret; ++} ++ ++layout(buffer_reference, std430, buffer_reference_align = 4) buffer decodeBufQ2_K { ++ block_q2_K block; ++}; ++ ++float16_t dequantFuncQ2_K(const in decodeBufQ2_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ const f16vec2 d = bl.block.d; ++ const uint idx = coordInBlock[1]; ++ const uint iqs = idx; ++ ++ const uint qsi = (iqs / 128) * 32 + (iqs % 32); // 0..31 ++ const uint scalesi = iqs / 16; // 0..15 ++ const uint qsshift = ((iqs % 128) / 32) * 2; // 0,2,4,6 ++ ++ uint32_t qs = bl.block.qs[qsi]; ++ const uint scales = bl.block.scales[scalesi]; ++ float16_t ret = d.x * float16_t(scales & 0xF) * float16_t((qs >> qsshift) & 3) - d.y * float16_t(scales >> 4); ++ return ret; ++} ++ ++layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ3_K { ++ block_q3_K block; ++}; ++ ++float16_t dequantFuncQ3_K(const in decodeBufQ3_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ const uint idx = coordInBlock[1]; ++ const uint iqs = idx; ++ ++ const uint n = iqs / 128; // 0,1 ++ const uint qsi = n * 32 + (iqs % 32); // 0..63 ++ const uint hmi = (iqs % 32); // 0..31 ++ const uint j = (iqs % 128) / 8; // 0..15 ++ const uint is = iqs / 16; // 0..15 ++ const uint halfsplit = ((iqs % 128) / 32); // 0,1,2,3 ++ const uint qsshift = halfsplit * 2; // 0,2,4,6 ++ const uint m = 1 << (4 * n + halfsplit); // 1,2,4,8,16,32,64,128 ++ ++ uint32_t scaleidx0 = (is < 8) ? is : (is-8); ++ uint32_t scaleidx0shift = (is < 8) ? 0 : 4; ++ uint32_t scaleidx1 = is + 8 - (is/4)*4; ++ uint32_t scaleidx1shift = (is/4)*2; ++ ++ const int8_t us = int8_t(((bl.block.scales[scaleidx0] >> scaleidx0shift) & 0xF) | (((bl.block.scales[scaleidx1] >> scaleidx1shift) & 3) << 4)); ++ ++ const float16_t dl = bl.block.d * float16_t(us - 32); ++ ++ float16_t ret = dl * float16_t(int8_t((bl.block.qs[qsi ] >> qsshift) & 3) - (((bl.block.hmask[hmi ] & m) != 0) ? 0 : 4)); ++ ++ return ret; ++} ++ ++layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ4_K { ++ block_q4_K block; ++}; ++ ++layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ4_K_packed16 { ++ block_q4_K_packed16 block; ++}; ++ ++float16_t dequantFuncQ4_K(const in decodeBufQ4_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ decodeBufQ4_K_packed16 bl16 = decodeBufQ4_K_packed16(bl); ++ const uint idx = coordInBlock[1]; ++ ++ const uint b = (idx & 0x20) >> 5; // 0,1 ++ const uint is = (idx & 0xE0) >> 5; // 0..7 ++ ++ const f16vec2 loadd = bl.block.d; ++ ++ uint32_t sc; ++ uint32_t mbyte; ++ ++ uint32_t scidx0 = (is < 4) ? is : (is + 4); ++ uint32_t scidx1 = (is < 4) ? is : (is - 4); ++ uint32_t scidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ uint32_t scidxshift1 = (is < 4) ? 0 : 2; ++ uint32_t mbidx0 = is + 4; ++ uint32_t mbidx1 = (is < 4) ? is + 4 : is; ++ uint32_t mbidxmask0 = (is < 4) ? 0xF : 0xF0; ++ uint32_t mbidxshift0 = (is < 4) ? 0 : 4; ++ uint32_t mbidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ uint32_t mbidxshift1 = (is < 4) ? 0 : 2; ++ ++ sc = uint8_t((bl.block.scales[scidx0] & 0xF) | ((bl.block.scales[scidx1] & scidxmask1) >> scidxshift1)); ++ mbyte = uint8_t(((bl.block.scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((bl.block.scales[mbidx1] & mbidxmask1) >> mbidxshift1)); ++ ++ const float16_t d = loadd.x * float16_t(sc); ++ const float16_t m = loadd.y * float16_t(mbyte); ++ ++ uint qs = uint32_t(bl16.block.qs[((idx & 0xC0) >> 2) + ((idx & 0x1E) >> 1)]); ++ qs = (qs >> (b * 4)) & 0x0F0F; ++ qs = unpack8(qs)[idx & 1]; ++ ++ float16_t ret = d * float16_t(qs) - m; ++ ++ return ret; ++} ++ ++layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ5_K { ++ block_q5_K block; ++}; ++ ++layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ5_K_packed16 { ++ block_q5_K_packed16 block; ++}; ++ ++float16_t dequantFuncQ5_K(const in decodeBufQ5_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ decodeBufQ5_K_packed16 bl16 = decodeBufQ5_K_packed16(bl); ++ const uint idx = coordInBlock[1]; ++ ++ const uint b = (idx & 0x20) >> 5; // 0,1 ++ const uint is = (idx & 0xE0) >> 5; // 0..7 ++ ++ const uint32_t hm = 0x0101 << is; ++ ++ const f16vec2 loadd = bl.block.d; ++ ++ uint32_t sc; ++ uint32_t mbyte; ++ ++ uint32_t scidx0 = (is < 4) ? is : (is + 4); ++ uint32_t scidx1 = (is < 4) ? is : (is - 4); ++ uint32_t scidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ uint32_t scidxshift1 = (is < 4) ? 0 : 2; ++ uint32_t mbidx0 = is + 4; ++ uint32_t mbidx1 = (is < 4) ? is + 4 : is; ++ uint32_t mbidxmask0 = (is < 4) ? 0xF : 0xF0; ++ uint32_t mbidxshift0 = (is < 4) ? 0 : 4; ++ uint32_t mbidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ uint32_t mbidxshift1 = (is < 4) ? 0 : 2; ++ ++ sc = uint8_t((bl.block.scales[scidx0] & 0xF) | ((bl.block.scales[scidx1] & scidxmask1) >> scidxshift1)); ++ mbyte = uint8_t(((bl.block.scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((bl.block.scales[mbidx1] & mbidxmask1) >> mbidxshift1)); ++ ++ const float16_t d = loadd.x * float16_t(sc); ++ const float16_t m = loadd.y * float16_t(mbyte); ++ ++ uint qh = uint32_t(bl16.block.qh[(idx & 0x1E) >> 1]); ++ qh = qh & hm; ++ qh = unpack8(qh)[idx & 1]; ++ ++ uint qs = uint32_t(bl16.block.qs[((idx & 0xC0) >> 2) + ((idx & 0x1E) >> 1)]); ++ qs = (qs >> (b * 4)) & 0x0F0F; ++ qs = unpack8(qs)[idx & 1]; ++ ++ float16_t ret = d * (float16_t(qs) + (qh != 0 ? float16_t(16) : float16_t(0))) - m; ++ ++ return ret; ++} ++ ++layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ6_K { ++ block_q6_K block; ++}; ++ ++layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ6_K_packed16 { ++ block_q6_K_packed16 block; ++}; ++ ++float16_t dequantFuncQ6_K(const in decodeBufQ6_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ decodeBufQ6_K_packed16 bl16 = decodeBufQ6_K_packed16(bl); ++ const uint idx = coordInBlock[1]; ++ ++ const uint b = (idx & 0x40) >> 6; // 0,1 ++ const uint qhshift = (idx & 0x60) >> 4; // 0,2,4,6 ++ const uint is = (idx & 0xF0) >> 4; // 0..15 ++ ++ const float16_t dscale = bl.block.d * float16_t(bl.block.scales[is]); ++ ++ uint ql = uint32_t(bl16.block.ql[((idx & 0x80) >> 2) + ((idx & 0x3E) >> 1)]); ++ ql = (ql >> (b * 4)) & 0x0F0F; ++ ++ uint qh = uint32_t(bl16.block.qh[((idx & 0x80) >> 3) + ((idx & 0x1E) >> 1)]); ++ qh = ((qh >> qhshift) & 0x0303) << 4; ++ ++ int q = unpack8(ql | qh)[idx & 1]; ++ ++ float16_t ret = dscale * float16_t(q - 32); ++ ++ return ret; ++} ++ ++#if defined(DATA_A_IQ4_NL) ++layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufIQ4_NL { ++ block_iq4_nl block; ++}; ++ ++float16_t dequantFuncIQ4_NL(const in decodeBufIQ4_NL bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ const float16_t d = bl.block.d; ++ const uint idx = coordInBlock[1]; ++ const uint iqs = idx & 0xF; ++ const uint shift = (idx & 0x10) >> 2; ++ uint32_t qs = bl.block.qs[iqs]; ++ qs >>= shift; ++ qs &= 0xF; ++ float16_t ret = float16_t(kvalues_iq4nl[qs]) * d; ++ return ret; ++} ++#endif ++ ++#if defined(DATA_A_Q4_0) ++#define dequantFuncA dequantFuncQ4_0 ++#elif defined(DATA_A_Q4_1) ++#define dequantFuncA dequantFuncQ4_1 ++#elif defined(DATA_A_Q5_0) ++#define dequantFuncA dequantFuncQ5_0 ++#elif defined(DATA_A_Q5_1) ++#define dequantFuncA dequantFuncQ5_1 ++#elif defined(DATA_A_Q8_0) ++#define dequantFuncA dequantFuncQ8_0 ++#elif defined(DATA_A_Q2_K) ++#define dequantFuncA dequantFuncQ2_K ++#elif defined(DATA_A_Q3_K) ++#define dequantFuncA dequantFuncQ3_K ++#elif defined(DATA_A_Q4_K) ++#define dequantFuncA dequantFuncQ4_K ++#elif defined(DATA_A_Q5_K) ++#define dequantFuncA dequantFuncQ5_K ++#elif defined(DATA_A_Q6_K) ++#define dequantFuncA dequantFuncQ6_K ++#elif defined(DATA_A_IQ4_NL) ++#define dequantFuncA dequantFuncIQ4_NL ++#endif +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp +new file mode 100644 +index 00000000..8d806435 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp +@@ -0,0 +1,13 @@ ++#extension GL_EXT_control_flow_attributes : require ++#extension GL_EXT_shader_16bit_storage : require ++ ++layout (push_constant) uniform parameter ++{ ++ uint M; ++ uint K; ++ uint stride_a; ++ uint stride_b; ++ uint nel; ++} p; ++ ++#include "types.comp" +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp +new file mode 100644 +index 00000000..8de14fc0 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp +@@ -0,0 +1,32 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {block_iq4_nl data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; ++ ++ init_iq4nl_shmem(); ++ ++ const uint tid = gl_LocalInvocationID.x % 64; ++ const uint il = tid/32; ++ const uint ir = tid%32; ++ const uint ib = 32*i + ir; ++ if (ib >= p.nel / 32) { ++ return; ++ } ++ ++ const uint q_idx = 8*il; ++ const uint b_idx = 1024*i + 32*ir + q_idx; ++ ++ const float d = float(data_a[ib].d); ++ ++ [[unroll]] for (uint l = 0; l < 8; ++l) { ++ data_b[b_idx + l + 0] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] & 0xF]); ++ data_b[b_idx + l + 16] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] >> 4]); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp +new file mode 100644 +index 00000000..157154af +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp +@@ -0,0 +1,34 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { ++ const uint i = gl_WorkGroupID.x * 256 + wgy; ++ if (i >= p.M * p.K / QUANT_K) { ++ return; ++ } ++ ++ const uint tid = gl_LocalInvocationID.x; ++ const uint ip = tid / 32; ++ const uint il = tid - 32 * ip; ++ const uint is = 8 * ip + il / 16; ++ ++ const uint y_idx = i * QUANT_K + 128 * ip + il; ++ ++ const uint ql_idx = 32 * ip + il; ++ const uint8_t qs = data_a[i].qs[32 * ip + il]; ++ ++ FLOAT_TYPE dall = FLOAT_TYPE(data_a[i].d.x); ++ FLOAT_TYPE dmin = FLOAT_TYPE(data_a[i].d.y); ++ data_b[y_idx + 0] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+0] & 0xF) * ((qs >> 0) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+0] >> 4)); ++ data_b[y_idx + 32] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+2] & 0xF) * ((qs >> 2) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+2] >> 4)); ++ data_b[y_idx + 64] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+4] & 0xF) * ((qs >> 4) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+4] >> 4)); ++ data_b[y_idx + 96] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+6] & 0xF) * ((qs >> 6) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+6] >> 4)); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp +new file mode 100644 +index 00000000..c17dd0d9 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp +@@ -0,0 +1,42 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { ++ const uint i = uint(gl_WorkGroupID.x * 256 + wgy); ++ if (i >= p.M * p.K / QUANT_K) { ++ return; ++ } ++ ++ const uint r = gl_LocalInvocationID.x / 4; ++ const uint tid = r / 2; ++ const uint is0 = r % 2; ++ const uint l0 = 16 * is0 + 4 * (gl_LocalInvocationID.x % 4); ++ const uint n = tid / 4; ++ const uint j = tid - 4*n; ++ ++ const uint8_t m = uint8_t(1 << (4*n + j)); ++ const uint is = 8*n + 2*j + is0; ++ const uint shift = 2*j; ++ ++ const int8_t us = int8_t(is < 4 ? (data_a[i].scales[is-0] & 0xF) | (((data_a[i].scales[is+8] >> 0) & 3) << 4) : ++ is < 8 ? (data_a[i].scales[is-0] & 0xF) | (((data_a[i].scales[is+4] >> 2) & 3) << 4) : ++ is < 12 ? (data_a[i].scales[is-8] >> 4) | (((data_a[i].scales[is+0] >> 4) & 3) << 4) : ++ (data_a[i].scales[is-8] >> 4) | (((data_a[i].scales[is-4] >> 6) & 3) << 4)); ++ const FLOAT_TYPE d_all = FLOAT_TYPE(data_a[i].d); ++ const FLOAT_TYPE dl = d_all * FLOAT_TYPE(us - 32); ++ ++ const uint y_idx = i * QUANT_K + 128 * n + 32 * j; ++ const uint qs_idx = 32*n; ++ ++ for (uint l = l0; l < l0 + 4; ++l) { ++ data_b[y_idx + l] = D_TYPE(dl * FLOAT_TYPE(int8_t((data_a[i].qs[qs_idx + l] >> shift) & 3) - (((data_a[i].hmask[l] & m) != 0) ? 0 : 4))); ++ } ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp +new file mode 100644 +index 00000000..40818532 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp +@@ -0,0 +1,30 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {block_q4_0 data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; ++ ++ const uint tid = gl_LocalInvocationID.x % 64; ++ const uint il = tid/32; ++ const uint ir = tid%32; ++ const uint ib = 32*i + ir; ++ if (ib >= p.nel / 32) { ++ return; ++ } ++ ++ const uint q_idx = 8*il; ++ const uint b_idx = 1024*i + 32*ir + q_idx; ++ ++ const float d = float(data_a[ib].d); ++ ++ [[unroll]] for (uint l = 0; l < 8; ++l) { ++ data_b[b_idx + l + 0] = D_TYPE(d * ((data_a[ib].qs[q_idx + l] & 0xF) - 8.0f)); ++ data_b[b_idx + l + 16] = D_TYPE(d * ((data_a[ib].qs[q_idx + l] >> 4) - 8.0f)); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp +new file mode 100644 +index 00000000..2f27eee6 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp +@@ -0,0 +1,32 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {block_q4_1 data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; ++ ++ const uint tid = gl_LocalInvocationID.x % 64; ++ const uint il = tid/32; ++ const uint ir = tid%32; ++ const uint ib = 32*i + ir; ++ if (ib >= p.nel / 32) { ++ return; ++ } ++ ++ const uint b_idx = 1024*i + 32*ir + 8*il; ++ ++ const float d = float(data_a[ib].d); ++ const float m = float(data_a[ib].m); ++ ++ const uint q_idx = 8*il; ++ ++ [[unroll]] for (uint l = 0; l < 8; ++l) { ++ data_b[b_idx + l + 0] = D_TYPE(d * (data_a[ib].qs[q_idx + l] & 0xF) + m); ++ data_b[b_idx + l + 16] = D_TYPE(d * (data_a[ib].qs[q_idx + l] >> 4) + m); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp +new file mode 100644 +index 00000000..987f113a +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp +@@ -0,0 +1,68 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { ++ const uint ib = gl_WorkGroupID.x * 256 + wgy; ++ if (ib >= p.M * p.K / QUANT_K) { ++ return; ++ } ++ ++ const uint tid = gl_LocalInvocationID.x; ++ const uint il = tid / 8; ++ const uint ir = tid % 8; ++ const uint is = 2 * il; ++ const uint n = 4; ++ ++ const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib].d.x); ++ const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib].d.y); ++ ++ const uint y_idx = ib * QUANT_K + 64 * il + n * ir; ++ const uint qs_idx = 32*il + n * ir; ++ ++ uint scidx0 = (is < 4) ? is : (is + 4); ++ uint scidx1 = (is < 4) ? is : (is - 4); ++ uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ uint scidxshift1 = (is < 4) ? 0 : 2; ++ uint mbidx0 = is + 4; ++ uint mbidx1 = (is < 4) ? is + 4 : is; ++ uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; ++ uint mbidxshift0 = (is < 4) ? 0 : 4; ++ uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ uint mbidxshift1 = (is < 4) ? 0 : 2; ++ ++ uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); ++ uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); ++ ++ const FLOAT_TYPE d1 = dall * sc; ++ const FLOAT_TYPE m1 = dmin * mbyte; ++ ++ scidx0 = (is < 4) ? is + 1 : (is + 5); ++ scidx1 = (is < 4) ? is + 1 : (is - 3); ++ scidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ scidxshift1 = (is < 4) ? 0 : 2; ++ mbidx0 = is + 5; ++ mbidx1 = (is < 4) ? is + 5 : is + 1; ++ mbidxmask0 = (is < 4) ? 0xF : 0xF0; ++ mbidxshift0 = (is < 4) ? 0 : 4; ++ mbidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ mbidxshift1 = (is < 4) ? 0 : 2; ++ ++ sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); ++ mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); ++ ++ const FLOAT_TYPE d2 = dall * sc; ++ const FLOAT_TYPE m2 = dmin * mbyte; ++ ++ [[unroll]] for (uint l = 0; l < n; ++l) { ++ data_b[y_idx + l ] = D_TYPE(d1 * FLOAT_TYPE(data_a[ib].qs[qs_idx + l] & 0xF) - m1); ++ data_b[y_idx + l + 32] = D_TYPE(d2 * FLOAT_TYPE(data_a[ib].qs[qs_idx + l] >> 4) - m2); ++ } ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp +new file mode 100644 +index 00000000..b20b8052 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp +@@ -0,0 +1,34 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {block_q5_0 data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; ++ ++ const uint tid = gl_LocalInvocationID.x % 64; ++ const uint il = tid/32; ++ const uint ir = tid%32; ++ const uint ib = 32*i + ir; ++ if (ib >= p.nel / 32) { ++ return; ++ } ++ ++ const uint b_idx = 1024*i + 32*ir + 8*il; ++ ++ const float d = float(data_a[ib].d); ++ const uint qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0]; ++ ++ const uint q_idx = 8*il; ++ ++ [[unroll]] for (uint l = 0; l < 8; ++l) { ++ const uint iqs = q_idx + l; ++ const uint vui = uint(data_a[ib].qs[iqs]); ++ data_b[b_idx + l + 0] = D_TYPE(d * (((vui & 0xF) | (((qh >> iqs) << 4) & 0x10)) - 16.0f)); ++ data_b[b_idx + l + 16] = D_TYPE(d * (((vui >> 4) | ((qh >> (iqs + 12)) & 0x10)) - 16.0f)); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp +new file mode 100644 +index 00000000..dc59fe3b +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp +@@ -0,0 +1,35 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {block_q5_1 data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; ++ ++ const uint tid = gl_LocalInvocationID.x % 64; ++ const uint il = tid/32; ++ const uint ir = tid%32; ++ const uint ib = 32*i + ir; ++ if (ib >= p.nel / 32) { ++ return; ++ } ++ ++ const uint b_idx = 1024*i + 32*ir + 8*il; ++ ++ const float d = float(data_a[ib].d); ++ const float m = float(data_a[ib].m); ++ const uint qh = data_a[ib].qh; ++ ++ const uint q_idx = 8*il; ++ ++ [[unroll]] for (uint l = 0; l < 8; ++l) { ++ const uint iqs = q_idx + l; ++ const uint vui = uint(data_a[ib].qs[iqs]); ++ data_b[b_idx + l + 0] = D_TYPE(d * (((vui & 0xF) | (((qh >> iqs) << 4) & 0x10))) + m); ++ data_b[b_idx + l + 16] = D_TYPE(d * (((vui >> 4) | ((qh >> (iqs + 12)) & 0x10))) + m); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp +new file mode 100644 +index 00000000..6db5403b +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp +@@ -0,0 +1,70 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { ++ const uint ib = gl_WorkGroupID.x * 256 + wgy; ++ if (ib >= p.M * p.K / QUANT_K) { ++ return; ++ } ++ ++ const uint tid = gl_LocalInvocationID.x; ++ const uint il = tid / 16; ++ const uint ir = tid % 16; ++ const uint is = 2 * il; ++ ++ const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib].d.x); ++ const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib].d.y); ++ ++ const uint y_idx = ib * QUANT_K + 64 * il + 2 * ir; ++ const uint qs_idx = 32*il + 2 * ir; ++ const uint qh_idx = 2 * ir; ++ ++ uint scidx0 = (is < 4) ? is : (is + 4); ++ uint scidx1 = (is < 4) ? is : (is - 4); ++ uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ uint scidxshift1 = (is < 4) ? 0 : 2; ++ uint mbidx0 = is + 4; ++ uint mbidx1 = (is < 4) ? is + 4 : is; ++ uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; ++ uint mbidxshift0 = (is < 4) ? 0 : 4; ++ uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ uint mbidxshift1 = (is < 4) ? 0 : 2; ++ ++ uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); ++ uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); ++ ++ const FLOAT_TYPE d1 = dall * sc; ++ const FLOAT_TYPE m1 = dmin * mbyte; ++ ++ scidx0 = (is < 4) ? is + 1 : (is + 5); ++ scidx1 = (is < 4) ? is + 1 : (is - 3); ++ scidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ scidxshift1 = (is < 4) ? 0 : 2; ++ mbidx0 = is + 5; ++ mbidx1 = (is < 4) ? is + 5 : is + 1; ++ mbidxmask0 = (is < 4) ? 0xF : 0xF0; ++ mbidxshift0 = (is < 4) ? 0 : 4; ++ mbidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ mbidxshift1 = (is < 4) ? 0 : 2; ++ ++ sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); ++ mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); ++ ++ const FLOAT_TYPE d2 = dall * sc; ++ const FLOAT_TYPE m2 = dmin * mbyte; ++ ++ const uint8_t hm1 = uint8_t(1 << (2 * il )); ++ const uint8_t hm2 = uint8_t(1 << (2 * il + 1)); ++ data_b[y_idx ] = D_TYPE(d1 * FLOAT_TYPE((data_a[ib].qs[qs_idx ] & 0xF) + (((data_a[ib].qh[qh_idx ] & hm1) != 0) ? 16 : 0)) - m1); ++ data_b[y_idx + 1] = D_TYPE(d1 * FLOAT_TYPE((data_a[ib].qs[qs_idx + 1] & 0xF) + (((data_a[ib].qh[qh_idx + 1] & hm1) != 0) ? 16 : 0)) - m1); ++ data_b[y_idx + 32] = D_TYPE(d2 * FLOAT_TYPE((data_a[ib].qs[qs_idx ] >> 4) + (((data_a[ib].qh[qh_idx ] & hm2) != 0) ? 16 : 0)) - m2); ++ data_b[y_idx + 33] = D_TYPE(d2 * FLOAT_TYPE((data_a[ib].qs[qs_idx + 1] >> 4) + (((data_a[ib].qh[qh_idx + 1] & hm2) != 0) ? 16 : 0)) - m2); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp +new file mode 100644 +index 00000000..0b913175 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp +@@ -0,0 +1,33 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { ++ const uint i = gl_WorkGroupID.x * 256 + wgy; ++ if (i >= p.M * p.K / QUANT_K) { ++ return; ++ } ++ const uint tid = gl_LocalInvocationID.x; ++ const uint ip = tid / 32; ++ const uint il = tid - 32 * ip; ++ const uint is = 8 * ip + il / 16; ++ ++ const uint y_idx = i * QUANT_K + 128 * ip + il; ++ ++ const uint ql_idx = 64 * ip + il; ++ const uint8_t qh = data_a[i].qh[32 * ip + il]; ++ ++ const FLOAT_TYPE d = FLOAT_TYPE(data_a[i].d); ++ ++ data_b[y_idx + 0] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 0] * (int8_t((data_a[i].ql[ql_idx + 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32))); ++ data_b[y_idx + 32] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 2] * (int8_t((data_a[i].ql[ql_idx + 32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32))); ++ data_b[y_idx + 64] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 4] * (int8_t((data_a[i].ql[ql_idx + 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32))); ++ data_b[y_idx + 96] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 6] * (int8_t((data_a[i].ql[ql_idx + 32] >> 4) | (((qh >> 6) & 3) << 4)) - 32))); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp +new file mode 100644 +index 00000000..bd1344a8 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp +@@ -0,0 +1,31 @@ ++#version 450 ++ ++#include "dequant_head.comp" ++ ++layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {block_q8_0 data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; ++ ++void main() { ++ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; ++ ++ const uint tid = gl_LocalInvocationID.x % 64; ++ const uint il = tid/32; ++ const uint ir = tid%32; ++ const uint ib = 32*i + ir; ++ if (ib >= p.nel / 32) { ++ return; ++ } ++ ++ const uint b_idx = 1024*i + 32*ir + 16*il; ++ ++ const float d = float(data_a[ib].d); ++ ++ const uint q_idx = 16*il; ++ ++ [[unroll]] for (uint l = 0; l < 16; l += 2) { ++ data_b[b_idx + l ] = D_TYPE(d * data_a[ib].qs[q_idx + l ]); ++ data_b[b_idx + l + 1] = D_TYPE(d * data_a[ib].qs[q_idx + l + 1]); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp +new file mode 100644 +index 00000000..4e68742b +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp +@@ -0,0 +1,34 @@ ++#version 450 ++ ++#extension GL_EXT_shader_16bit_storage : require ++#extension GL_EXT_control_flow_attributes : enable ++ ++layout (push_constant) uniform parameter ++{ ++ uint ncols; ++ uint rows_per_channel; ++ uint n_past; ++} p; ++ ++#include "types.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++void main() { ++ const uint col = gl_GlobalInvocationID.y; ++ const uint row = gl_GlobalInvocationID.x; ++ ++ if (col >= p.ncols) { ++ return; ++ } ++ ++ const uint i = row*p.ncols + col; ++ if (col > p.n_past + row % p.rows_per_channel) { ++ data_d[i] = D_TYPE(uintBitsToFloat(0xFF800000)); ++ } else { ++ data_d[i] = D_TYPE(data_a[i]); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp +new file mode 100644 +index 00000000..9fb69c6c +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp +@@ -0,0 +1,27 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_binary_head.comp" ++ ++const uint num_threads = 256; ++ ++layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ uint idx = get_idx(); ++ ++ // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation ++ const uint num_iter = 2; ++ ++ [[unroll]] for (uint i = 0; i < num_iter; ++i) { ++ if (idx >= p.ne) { ++ continue; ++ } ++ uint i00, i01, i02, i03; ++ get_indices(idx, i00, i01, i02, i03); ++ ++ data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) / FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); ++ ++ idx += num_threads; ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp +new file mode 100644 +index 00000000..c5be8131 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp +@@ -0,0 +1,289 @@ ++#version 450 ++ ++#extension GL_EXT_control_flow_attributes : enable ++#extension GL_EXT_shader_16bit_storage : require ++ ++#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require ++#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require ++#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require ++#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require ++ ++#extension GL_KHR_memory_scope_semantics : enable ++#extension GL_KHR_cooperative_matrix : enable ++#extension GL_NV_cooperative_matrix2 : enable ++#extension GL_EXT_buffer_reference : enable ++#extension GL_KHR_shader_subgroup_ballot : enable ++#extension GL_KHR_shader_subgroup_vote : enable ++#extension GL_EXT_null_initializer : enable ++ ++#include "types.comp" ++#include "dequant_funcs_cm2.comp" ++ ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++layout (constant_id = 1) const uint32_t Br = 32; ++layout (constant_id = 2) const uint32_t Bc = 32; ++layout (constant_id = 3) const uint32_t D = 32; ++layout (constant_id = 4) const uint32_t Clamp = gl_CooperativeMatrixClampModeConstantNV; ++ ++layout (push_constant) uniform parameter { ++ uint32_t N; ++ uint32_t KV; ++ ++ uint32_t ne1; ++ uint32_t ne2; ++ uint32_t ne3; ++ ++ uint32_t neq2; ++ uint32_t neq3; ++ uint32_t nek2; ++ uint32_t nek3; ++ uint32_t nev2; ++ uint32_t nev3; ++ uint32_t nem1; ++ ++ uint32_t nb02; ++ uint32_t nb03; ++ uint32_t nb12; ++ uint32_t nb13; ++ uint32_t nb22; ++ uint32_t nb23; ++ uint32_t nb31; ++ ++ float scale; ++ float max_bias; ++ float logit_softcap; ++ ++ uint32_t mask; ++ uint32_t n_head_log2; ++ float m0; ++ float m1; ++} p; ++ ++layout (binding = 0) readonly buffer Q {uint8_t data_q[];}; ++layout (binding = 1) readonly buffer K {uint8_t data_k[];}; ++layout (binding = 2) readonly buffer V {uint8_t data_v[];}; ++layout (binding = 3) readonly buffer M {uint8_t data_m[];}; ++layout (binding = 4) writeonly buffer O {D_TYPE data_o[];}; ++ ++#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b)) ++ ++ACC_TYPE maxReduce(const in ACC_TYPE x, const in ACC_TYPE y) { ++ return max(x, y); ++} ++ ++ACC_TYPE smearReduce(const in ACC_TYPE x, const in ACC_TYPE y) { ++ return x; ++} ++ ++// Replace matrix elements >= numRows or numCols with 'replace' ++ACC_TYPE replacePadding(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem, const in ACC_TYPE replace, const in uint32_t numRows, const in uint32_t numCols) { ++ if (row >= numRows || col >= numCols) { ++ return replace; ++ } ++ return elem; ++} ++ ++ACC_TYPE Exp(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem) ++{ ++ return exp(elem); ++} ++ ++ACC_TYPE Max(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem0, const in ACC_TYPE elem1) ++{ ++ return max(elem0, elem1); ++} ++ ++#if defined(BLOCK_SIZE) ++#define DECODEFUNC , DEQUANTFUNC ++#else ++#define DECODEFUNC ++#endif ++ ++void main() { ++#if defined(DATA_A_IQ4_NL) ++ init_iq4nl_shmem(); ++#endif ++ ++ const uint32_t N = p.N; ++ const uint32_t KV = p.KV; ++ ++ const uint32_t Tr = CEIL_DIV(N, Br); ++ const uint32_t Tc = CEIL_DIV(KV, Bc); ++ ++ const uint32_t i = gl_WorkGroupID.x; ++ ++ const uint32_t iq2 = gl_WorkGroupID.y; ++ const uint32_t iq3 = gl_WorkGroupID.z; ++ ++ // broadcast factors ++ const uint32_t rk2 = p.neq2/p.nek2; ++ const uint32_t rk3 = p.neq3/p.nek3; ++ ++ const uint32_t rv2 = p.neq2/p.nev2; ++ const uint32_t rv3 = p.neq3/p.nev3; ++ ++ // k indices ++ const uint32_t ik3 = iq3 / rk3; ++ const uint32_t ik2 = iq2 / rk2; ++ ++ // v indices ++ const uint32_t iv3 = iq3 / rv3; ++ const uint32_t iv2 = iq2 / rv2; ++ ++ tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutQ = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); ++ tensorLayoutNV<2, Clamp> tensorLayoutK = createTensorLayoutNV(2, Clamp); ++ tensorLayoutNV<2, Clamp> tensorLayoutV = createTensorLayoutNV(2, Clamp); ++ ++ tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0); ++ ++#if defined(BLOCK_SIZE) ++ tensorLayoutK = setTensorLayoutBlockSizeNV(tensorLayoutK, 1, BLOCK_SIZE); ++ tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, BLOCK_SIZE); ++#endif ++ ++ tensorLayoutQ = setTensorLayoutDimensionNV(tensorLayoutQ, N, D); ++ tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, KV, D); ++ tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, KV, D); ++ ++ coopmat Q; ++ coopmat Qf16; ++ ++ uint32_t q_offset = iq2*p.nb02+iq3*p.nb03; ++ coopMatLoadTensorNV(Q, data_q, q_offset, sliceTensorLayoutNV(tensorLayoutQ, i * Br, Br, 0, D)); ++ ++ Qf16 = coopmat(Q); ++ Qf16 *= float16_t(p.scale); ++ ++ coopmat O = coopmat(0); ++ ++ coopmat L, M; ++ ++ L = coopmat(0); ++ M = coopmat(-1.0/0.0); ++ ++ ACC_TYPE slope = ACC_TYPE(1.0); ++ ++ // ALiBi ++ if (p.max_bias > 0.0f) { ++ const uint32_t h = iq2; ++ ++ const ACC_TYPE base = ACC_TYPE(h < p.n_head_log2 ? p.m0 : p.m1); ++ const int exph = int(h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1); ++ ++ slope = pow(base, ACC_TYPE(exph)); ++ } ++ ++ [[dont_unroll]] ++ for (uint32_t j = 0; j < Tc; ++j) { ++ ++ coopmat S = coopmat(0); ++ ++ coopmat K_T; ++ ++ uint32_t k_offset = ik2*p.nb12 + ik3*p.nb13; ++ coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, D), tensorViewTranspose DECODEFUNC); ++ S = coopMatMulAdd(Qf16, K_T, S); ++ ++ if (p.logit_softcap != 0.0f) { ++ [[unroll]] ++ for (int k = 0; k < S.length(); ++k) { ++ S[k] = ACC_TYPE(p.logit_softcap)*tanh(S[k]); ++ } ++ } ++ ++ if (p.mask != 0) { ++ tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutM = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); ++ tensorLayoutM = setTensorLayoutDimensionNV(tensorLayoutM, p.nem1, KV); ++ ++ coopmat mv; ++ ++ coopMatLoadTensorNV(mv, data_m, 0, sliceTensorLayoutNV(tensorLayoutM, i * Br, Br, j * Bc, Bc)); ++ ++ S += slope*coopmat(mv); ++ } ++ ++ // Clear padding elements to -inf, so they don't contribute to rowmax ++ if (Clamp != 0 && ++ ((j + 1) * Bc > KV || ++ (i + 1) * Br > N)) { ++ ++ uint R = ((i + 1) * Br > N) ? (N % Br) : Br; ++ uint C = ((j + 1) * Bc > KV) ? (KV % Bc) : Bc; ++ ++ coopMatPerElementNV(S, S, replacePadding, ACC_TYPE(-1.0/0.0), R, C); ++ } ++ ++ coopmat rowmax, P, rowsum, eM; ++ ++ coopMatReduceNV(rowmax, S, gl_CooperativeMatrixReduceRowNV, maxReduce); ++ ++ coopmat Mold = M; ++ ++ // M = max(rowmax, Mold) ++ // P = e^(S - M) ++ // eM = e^(Mold - M) ++ coopMatPerElementNV(M, rowmax, Max, Mold); ++ coopMatPerElementNV(P, S - M, Exp); ++ coopMatPerElementNV(eM, Mold - M, Exp); ++ ++ // Clear padding elements to 0, so they don't contribute to rowsum ++ if (Clamp != 0 && ++ ((j + 1) * Bc > KV || ++ (i + 1) * Br > N)) { ++ ++ uint R = ((i + 1) * Br > N) ? (N % Br) : Br; ++ uint C = ((j + 1) * Bc > KV) ? (KV % Bc) : Bc; ++ ++ coopMatPerElementNV(P, P, replacePadding, ACC_TYPE(0.0), R, C); ++ } ++ ++ coopmat P_A = coopmat(P); ++ ++ // compute rowsum by multiplying by matrix of all ones. ++ coopmat One = coopmat(1.0); ++ ++ rowsum = coopmat(0.0); ++ rowsum = coopMatMulAdd(P_A, One, rowsum); ++ ++ coopmat V; ++ uint32_t v_offset = iv2*p.nb22 + iv3*p.nb23; ++ coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, D) DECODEFUNC); ++ ++ L = eM*L + rowsum; ++ ++ // This is the "diagonal" matrix in the paper, but since we do componentwise ++ // multiply rather than matrix multiply it has the diagonal element smeared ++ // across the row ++ coopmat eMdiag; ++ ++ // resize eM by using smear/reduce ++ coopMatReduceNV(eMdiag, eM, gl_CooperativeMatrixReduceRowNV, smearReduce); ++ ++ O = eMdiag * O; ++ ++ O = coopMatMulAdd(P_A, V, O); ++ } ++ ++ coopmat Ldiag; ++ ++ // resize L by using smear/reduce ++ coopMatReduceNV(Ldiag, L, gl_CooperativeMatrixReduceRowNV, smearReduce); ++ ++ [[unroll]] ++ for (int k = 0; k < Ldiag.length(); ++k) { ++ Ldiag[k] = ACC_TYPE(1.0) / Ldiag[k]; ++ } ++ ++ O = Ldiag*O; ++ ++ tensorLayoutNV<3, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(3, gl_CooperativeMatrixClampModeConstantNV); ++ tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.ne2, p.ne1, D); ++ ++ // permute dimensions ++ tensorViewNV<3, false, 1, 0, 2> tensorViewPermute = createTensorViewNV(3, false, 1, 0, 2); ++ uint32_t o_offset = iq3*p.ne2*p.ne1; ++ ++ coopmat O_D = coopmat(O); ++ coopMatStoreTensorNV(O_D, data_o, o_offset, sliceTensorLayoutNV(tensorLayoutD, i * Br, Br, iq2, 1, 0, D), tensorViewPermute); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp +new file mode 100644 +index 00000000..4cc7a68c +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp +@@ -0,0 +1,25 @@ ++#version 450 ++ ++#include "generic_head.comp" ++#include "types.comp" ++ ++#extension GL_EXT_control_flow_attributes : enable ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++void main() { ++ const float GELU_COEF_A = 0.044715f; ++ const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; ++ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; ++ ++ if (i >= p.KX) { ++ return; ++ } ++ ++ const float xi = float(data_a[i]); ++ const float val = SQRT_2_OVER_PI*xi*(1.0f + GELU_COEF_A*xi*xi); ++ data_d[i] = D_TYPE(0.5f*xi*(2.0f - 2.0f / (exp(2 * val) + 1))); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp +new file mode 100644 +index 00000000..e6e6fcfd +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp +@@ -0,0 +1,23 @@ ++#version 450 ++ ++#include "generic_head.comp" ++#include "types.comp" ++ ++#extension GL_EXT_control_flow_attributes : enable ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++void main() { ++ const float GELU_QUICK_COEF = -1.702f; ++ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; ++ ++ if (i >= p.KX) { ++ return; ++ } ++ ++ const float x = float(data_a[i]); ++ data_d[i] = D_TYPE(x * (1.0f / (1.0f + exp(GELU_QUICK_COEF * x)))); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp +new file mode 100644 +index 00000000..062e2a4c +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp +@@ -0,0 +1,64 @@ ++#extension GL_EXT_shader_16bit_storage : require ++#extension GL_EXT_control_flow_attributes : require ++ ++layout (push_constant) uniform parameter ++{ ++ uint ne; ++ uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03; ++ uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13; ++ uint ne20; uint ne21; uint ne22; uint ne23; uint nb20; uint nb21; uint nb22; uint nb23; ++ uint misalign_offsets; ++ float param1; float param2; int param3; ++} p; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; ++layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; ++ ++// true if src0/src1 are the same shape and the indices can be reused without additional modulus ++layout(constant_id = 0) const bool norepeat = false; ++ ++uint get_idx() { ++ return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; ++} ++ ++uint get_aoffset() { return p.misalign_offsets >> 16; } ++uint get_boffset() { return (p.misalign_offsets >> 8) & 0xFF; } ++uint get_doffset() { return p.misalign_offsets & 0xFF; } ++ ++// mod and div are expensive and coordinates/dimensions are often power of 2 or equal to 1 ++uint fastmod(uint a, uint b) { ++ if ((b & (b-1)) == 0) { ++ return a & (b-1); ++ } ++ return a % b; ++} ++ ++uint fastdiv(uint a, uint b) { ++ return (a < b) ? 0 : (a / b); ++} ++ ++void get_indices(uint idx, out uint i00, out uint i01, out uint i02, out uint i03) { ++ i03 = fastdiv(idx, (p.ne02*p.ne01*p.ne00)); ++ const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00; ++ i02 = fastdiv((idx - i03_offset), (p.ne01*p.ne00)); ++ const uint i02_offset = i02*p.ne01*p.ne00; ++ i01 = (idx - i03_offset - i02_offset) / p.ne00; ++ i00 = idx - i03_offset - i02_offset - i01*p.ne00; ++} ++ ++uint src0_idx(uint i00, uint i01, uint i02, uint i03) { ++ return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00; ++} ++ ++uint src1_idx(uint i00, uint i01, uint i02, uint i03) { ++ if (norepeat) { ++ return i03*p.nb13 + i02*p.nb12 + i01*p.nb11 + i00*p.nb10; ++ } else { ++ return fastmod(i03, p.ne13)*p.nb13 + fastmod(i02, p.ne12)*p.nb12 + fastmod(i01, p.ne11)*p.nb11 + fastmod(i00, p.ne10)*p.nb10; ++ } ++} ++ ++uint dst_idx(uint i00, uint i01, uint i02, uint i03) { ++ return i03*p.nb23 + i02*p.nb22 + i01*p.nb21 + i00*p.nb20; ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp +new file mode 100644 +index 00000000..66e46ae6 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp +@@ -0,0 +1,9 @@ ++#extension GL_EXT_shader_16bit_storage : require ++ ++layout (push_constant) uniform parameter ++{ ++ uint KX; ++ uint KY; ++ float param1; ++ float param2; ++} p; +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp +new file mode 100644 +index 00000000..68d1bc9f +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp +@@ -0,0 +1,56 @@ ++#extension GL_EXT_shader_16bit_storage : require ++#extension GL_EXT_control_flow_attributes : require ++ ++layout (push_constant) uniform parameter ++{ ++ uint ne; ++ uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03; ++ uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13; ++ uint misalign_offsets; ++ float param1; float param2; ++ ++ uint ne0_012mp; uint ne0_012L; ++ uint ne0_01mp; uint ne0_01L; ++ uint ne0_0mp; uint ne0_0L; ++ uint ne1_012mp; uint ne1_012L; ++ uint ne1_01mp; uint ne1_01L; ++ uint ne1_0mp; uint ne1_0L; ++} p; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++uint get_idx() { ++ return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; ++} ++ ++uint get_aoffset() { return p.misalign_offsets >> 16; } ++uint get_doffset() { return p.misalign_offsets & 0xFFFF; } ++ ++// see init_fastdiv_values in ggml-vulkan.cpp ++uint fastdiv(uint n, uint mp, uint L) { ++ uint msbs, lsbs; ++ // msbs = mulhi(n, mp) ++ umulExtended(n, mp, msbs, lsbs); ++ return (msbs + n) >> L; ++} ++ ++uint src0_idx(uint idx) { ++ const uint i03 = fastdiv(idx, p.ne0_012mp, p.ne0_012L); ++ const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00; ++ const uint i02 = fastdiv(idx - i03_offset, p.ne0_01mp, p.ne0_01L); ++ const uint i02_offset = i02*p.ne01*p.ne00; ++ const uint i01 = fastdiv(idx - i03_offset - i02_offset, p.ne0_0mp, p.ne0_0L); ++ const uint i00 = idx - i03_offset - i02_offset - i01*p.ne00; ++ return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00; ++} ++ ++uint dst_idx(uint idx) { ++ const uint i13 = fastdiv(idx, p.ne1_012mp, p.ne1_012L); ++ const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10; ++ const uint i12 = fastdiv(idx - i13_offset, p.ne1_01mp, p.ne1_01L); ++ const uint i12_offset = i12*p.ne11*p.ne10; ++ const uint i11 = fastdiv(idx - i13_offset - i12_offset, p.ne1_0mp, p.ne1_0L); ++ const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10; ++ return i13*p.nb13 + i12*p.nb12 + i11*p.nb11 + i10*p.nb10; ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp +new file mode 100644 +index 00000000..e877ed77 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp +@@ -0,0 +1,28 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_binary_head.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ const uint i00 = gl_GlobalInvocationID.x; ++ const uint i10 = gl_GlobalInvocationID.y; ++ const uint i11 = (gl_GlobalInvocationID.z)/p.ne12; ++ const uint i12 = (gl_GlobalInvocationID.z)%p.ne12; ++ ++ if (i00 >= p.ne00) { ++ return; ++ } ++ ++ const uint i01 = data_b[get_boffset() + i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; ++ ++ const uint a_offset = get_aoffset() + i01*p.nb01 + i11*p.nb02 + i12*p.nb03; ++ const uint d_offset = get_doffset() + i10*p.nb21 + i11*p.nb22 + i12*p.nb23; ++ ++#ifndef OPTIMIZATION_ERROR_WORKAROUND ++ data_d[d_offset + i00] = D_TYPE(data_a[a_offset + i00]); ++#else ++ data_d[d_offset + i00] = data_a[a_offset + i00]; ++#endif ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp +new file mode 100644 +index 00000000..1426fde6 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp +@@ -0,0 +1,39 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_binary_head.comp" ++#include "dequant_funcs.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ const uint i00 = (gl_GlobalInvocationID.x)*2; ++ const uint i10 = gl_GlobalInvocationID.y; ++ const uint i11 = (gl_GlobalInvocationID.z)/p.ne12; ++ const uint i12 = (gl_GlobalInvocationID.z)%p.ne12; ++ ++#if defined(DATA_A_IQ4_NL) ++ init_iq4nl_shmem(); ++#endif ++ ++ if (i00 >= p.ne00) { ++ return; ++ } ++ ++ const uint i01 = data_b[i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; ++ ++ const uint a_offset = i01*p.nb01 + i11*p.nb02 + i12*p.nb03; ++ const uint d_offset = i10*p.nb21 + i11*p.nb22 + i12*p.nb23; ++ ++ const uint ib = a_offset + i00/QUANT_K; // block index ++ const uint iqs = (i00%QUANT_K)/QUANT_R; // quant index ++ const uint iybs = i00 - i00%QUANT_K; // dst block start index ++ const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2; ++ ++ vec2 v = dequantize(ib, iqs, 0); ++ const vec2 dm = get_dm(ib, 0); ++ v = v * dm.x + dm.y; ++ ++ data_d[d_offset + iybs + iqs ] = D_TYPE(v.x); ++ data_d[d_offset + iybs + iqs + y_offset] = D_TYPE(v.y); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp +new file mode 100644 +index 00000000..b6a0d564 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp +@@ -0,0 +1,66 @@ ++#version 450 ++ ++#include "generic_head.comp" ++#include "types.comp" ++ ++#extension GL_EXT_control_flow_attributes : enable ++#define BLOCK_SIZE 512 ++ ++layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++shared float tmp[BLOCK_SIZE]; ++ ++void main() { ++ const uint group_size = p.KX; ++ const float eps = p.param1; ++ ++ const uint tid = gl_LocalInvocationID.x; ++ const uint start = gl_WorkGroupID.x * group_size + tid; ++ const uint end = (gl_WorkGroupID.x + 1) * group_size; ++ ++ tmp[tid] = 0.0f; ++ ++ // Calculate mean ++ [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { ++ tmp[tid] += float(data_a[col]); ++ } ++ ++ // tmp up partial tmps and write back result ++ barrier(); ++ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { ++ if (tid < s) { ++ tmp[tid] += tmp[tid + s]; ++ } ++ barrier(); ++ } ++ ++ const float mean = tmp[0] / group_size; ++ barrier(); ++ tmp[tid] = 0.0f; ++ ++ // Calculate variance ++ [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { ++ const float xi = float(data_a[col]) - mean; ++ data_d[col] = D_TYPE(xi); ++ tmp[tid] += xi * xi; ++ } ++ ++ // sum up partial sums and write back result ++ barrier(); ++ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { ++ if (tid < s) { ++ tmp[tid] += tmp[tid + s]; ++ } ++ barrier(); ++ } ++ ++ const float variance = tmp[0] / group_size; ++ const float scale = inversesqrt(variance + eps); ++ ++ [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { ++ data_d[col] *= D_TYPE(scale); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp +new file mode 100644 +index 00000000..122b1e93 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp +@@ -0,0 +1,87 @@ ++#version 450 ++ ++#extension GL_EXT_shader_16bit_storage : require ++#extension GL_EXT_spirv_intrinsics: enable ++#extension GL_EXT_control_flow_attributes : require ++ ++#if RTE16 ++spirv_execution_mode(capabilities = [4467], 4462, 16); // RoundingModeRTE, 16 bits ++#endif ++ ++layout (push_constant) uniform parameter ++{ ++ uint batch_offset; uint offset_delta; ++ uint IC; ++ uint IW; uint IH; ++ uint OW; uint OH; ++ uint KW; uint KH; ++ uint pelements; ++ uint CHW; ++ int s0; int s1; ++ int p0; int p1; ++ int d0; int d1; ++} p; ++ ++#include "types.comp" ++ ++layout(constant_id = 0) const uint BLOCK_SIZE = 32; ++ ++const uint NUM_ITER = 512 / BLOCK_SIZE; ++ ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++void main() { ++ const uint gidx = gl_GlobalInvocationID.x; ++ ++ const uint oh = gl_GlobalInvocationID.y; ++ const uint batch = gl_GlobalInvocationID.z / p.IC; ++ const uint ic = gl_GlobalInvocationID.z % p.IC; ++ ++ A_TYPE values[NUM_ITER]; ++ uint offset_dst[NUM_ITER]; ++ [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { ++ values[idx] = A_TYPE(0); ++ } ++ ++ [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { ++ ++ const uint i = gidx * NUM_ITER + idx; ++ ++ const uint ksize = p.OW * (p.KH > 1 ? p.KW : 1); ++ const uint kx = i / ksize; ++ const uint kd = kx * ksize; ++ const uint ky = (i - kd) / p.OW; ++ const uint ix = i % p.OW; ++ ++ const uint iiw = ix * p.s0 + kx * p.d0 - p.p0; ++ const uint iih = oh * p.s1 + ky * p.d1 - p.p1; ++ ++ offset_dst[idx] = ++ ((batch * p.OH + oh) * p.OW + ix) * p.CHW + ++ (ic * (p.KW * p.KH) + ky * p.KW + kx); ++ ++ if (i >= p.pelements) { ++ continue; ++ } ++ ++ if (iih < p.IH && iiw < p.IW) { ++ const uint offset_src = ic * p.offset_delta + batch * p.batch_offset; ++ values[idx] = data_a[offset_src + iih * p.IW + iiw]; ++ } ++ } ++ ++ [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { ++ ++ const uint i = gidx * NUM_ITER + idx; ++ ++ if (i >= p.pelements) { ++ continue; ++ } ++ ++ data_d[offset_dst[idx]] = D_TYPE(values[idx]); ++ } ++ ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp +new file mode 100644 +index 00000000..d90a99ae +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp +@@ -0,0 +1,22 @@ ++#version 450 ++ ++#include "generic_head.comp" ++#include "types.comp" ++ ++#extension GL_EXT_control_flow_attributes : enable ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++void main() { ++ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; ++ ++ if (i >= p.KX) { ++ return; ++ } ++ ++ const float val = float(data_a[i]); ++ data_d[i] = D_TYPE(max(val, 0.0f) + min(val, 0.0f) * p.param1); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp +new file mode 100644 +index 00000000..43de19df +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp +@@ -0,0 +1,27 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_binary_head.comp" ++ ++const uint num_threads = 256; ++ ++layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ uint idx = get_idx(); ++ ++ // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation ++ const uint num_iter = 2; ++ ++ [[unroll]] for (uint i = 0; i < num_iter; ++i) { ++ if (idx >= p.ne) { ++ continue; ++ } ++ uint i00, i01, i02, i03; ++ get_indices(idx, i00, i01, i02, i03); ++ ++ data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) * FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); ++ ++ idx += num_threads; ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp +new file mode 100644 +index 00000000..4c64fd47 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp +@@ -0,0 +1,48 @@ ++#version 450 ++ ++#extension GL_EXT_control_flow_attributes : enable ++ ++layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {float data_a[];}; ++layout (binding = 0) readonly buffer A4 {vec4 data_a4[];}; ++layout (binding = 1) writeonly buffer D {float data_d[];}; ++layout (binding = 1) writeonly buffer D4 {vec4 data_d4[];}; ++ ++layout (push_constant) uniform parameter { ++ uint ne; ++ uint k_num; ++} p; ++ ++void main() { ++ // Each invocation handles four consecutive components ++ const uint idx = gl_GlobalInvocationID.x * 4; ++ ++ if (idx >= p.ne) { ++ return; ++ } ++ ++ // Check if all four components are in bounds and aligned, ++ // then use vector loads ++ if (idx + 3 < p.ne && (p.ne % 4) == 0) { ++ vec4 result = vec4(0.0f); ++ ++ [[unroll]] for (uint i = 0; i < p.k_num; i++) { ++ result += data_a4[(i * p.ne + idx) / 4]; ++ } ++ ++ data_d4[idx / 4] = result; ++ } else { ++ [[unroll]] for (uint j = 0; j < 4; ++j) { ++ if (idx + j < p.ne) { ++ float result = 0.0f; ++ ++ [[unroll]] for (uint i = 0; i < p.k_num; i++) { ++ result += data_a[i * p.ne + idx + j]; ++ } ++ ++ data_d[idx + j] = result; ++ } ++ } ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp +new file mode 100644 +index 00000000..24875cdc +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp +@@ -0,0 +1,152 @@ ++#version 450 ++ ++#ifdef FLOAT16 ++#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require ++#endif ++#extension GL_EXT_shader_explicit_arithmetic_types : require ++ ++#include "mul_mat_vec_base.comp" ++ ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++#if !defined(DATA_A_F32) && !defined(DATA_A_F16) ++#define K_PER_ITER 8 ++#else ++#define K_PER_ITER 2 ++#endif ++ ++ ++uint a_offset, b_offset, d_offset, y_offset; ++ ++void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint tid, const uint i, bool lastiter) ++{ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ const uint col = i*BLOCK_SIZE + K_PER_ITER*tid; ++ const uint iqs = (col%QUANT_K)/QUANT_R; // quant index ++ const uint iybs = col - col%QUANT_K; // y block start index ++ ++#if K_PER_ITER == 8 ++#if QUANT_R == 2 ++ const B_TYPE_VEC4 bv02 = data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]; ++ const B_TYPE_VEC4 bv13 = data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs + y_offset) / 4]; ++ const vec4 bv0 = vec4(bv02.x, bv13.x, bv02.y, bv13.y); ++ const vec4 bv1 = vec4(bv02.z, bv13.z, bv02.w, bv13.w); ++#else ++ const vec4 bv0 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]); ++ const vec4 bv1 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4 + 1]); ++#endif ++#else ++ // Check if the second of the pair of elements is OOB, and don't fetch B or ++ // accumulate it. We still fetch a pair of elements for A, which is fine for ++ // quantized formats since they'll be within the same block. We should ++ // probably skip fetching the second element for F16/F32, but as of now we ++ // still do. ++ const bool OOB = lastiter && (iybs + iqs + y_offset >= p.ncols); ++ ++ FLOAT_TYPE b0 = 0, b1 = 0; ++ b0 = FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs]); ++ if (!OOB) { ++ b1 = FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs + y_offset]); ++ } ++#endif ++ uint ibi = first_row*p.ncols; ++ [[unroll]] for (uint n = 0; n < num_rows; ++n) { ++ const uint ib = (ibi + col)/QUANT_K; // block index ++ ibi += p.ncols; ++ ++#if K_PER_ITER == 8 ++ vec4 v = dequantize4(ib, iqs, a_offset); ++ vec4 v2 = dequantize4(ib, iqs+(4/QUANT_R), a_offset); ++ ++ const vec2 dm = get_dm(ib, a_offset); ++ if (dm.y != 0) { // quant has min component ++ v = v * dm.x + dm.y; ++ v2 = v2 * dm.x + dm.y; ++ } ++ ++ // matrix multiplication ++ FLOAT_TYPE rowtmp = dot(bv0, v); ++ rowtmp += dot(bv1, v2); ++ ++ if (dm.y == 0) ++ rowtmp *= dm.x; ++ ++ temp[j][n] += rowtmp; ++#else ++ const vec2 v = dequantize(ib, iqs, a_offset); ++ ++ // matrix multiplication ++ temp[j][n] = fma(FLOAT_TYPE(v.x), b0, temp[j][n]); ++ if (!OOB) { ++ temp[j][n] = fma(FLOAT_TYPE(v.y), b1, temp[j][n]); ++ } ++#endif ++ } ++ } ++} ++ ++void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { ++ const uint tid = gl_LocalInvocationID.x; ++ ++ get_offsets(a_offset, b_offset, d_offset); ++ a_offset /= QUANT_K; ++ ++ y_offset = QUANT_R == 1 ? 1 : QUANT_K/2; ++ ++ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; ++ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { ++ temp[j][i] = FLOAT_TYPE(0); ++ } ++ } ++ ++ uint num_iters = p.ncols / (K_PER_ITER * BLOCK_SIZE); ++ if (num_iters * K_PER_ITER * BLOCK_SIZE + K_PER_ITER*tid < p.ncols) { ++ num_iters++; ++ } ++ int unroll_count = 4; ++ uint unrolled_iters = num_iters & ~(unroll_count - 1); ++ ++ uint i = 0; ++ while (i < unrolled_iters) { ++ // Manually partially unroll the loop ++ [[unroll]] for (uint k = 0; k < unroll_count; ++k) { ++ iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false); ++ i++; ++ } ++ } ++ unroll_count = 2; ++ unrolled_iters = num_iters & ~(unroll_count - 1); ++ while (i < unrolled_iters) { ++ // Manually partially unroll the loop ++ [[unroll]] for (uint k = 0; k < unroll_count; ++k) { ++ iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false); ++ i++; ++ } ++ } ++ while (i < num_iters) { ++ iter(temp, first_row, num_rows, tid, i*K_PER_ITER, true); ++ i++; ++ } ++ ++ reduce_result(temp, d_offset, first_row, num_rows, tid); ++} ++ ++void main() { ++ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); ++ ++#if defined(DATA_A_IQ4_NL) ++ init_iq4nl_shmem(); ++#endif ++ ++ // do NUM_ROWS at a time, unless there aren't enough remaining rows ++ if (first_row + NUM_ROWS <= p.stride_d) { ++ compute_outputs(first_row, NUM_ROWS); ++ } else { ++ if (first_row >= p.stride_d) { ++ return; ++ } ++ compute_outputs(first_row, p.stride_d - first_row); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp +new file mode 100644 +index 00000000..903753c7 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp +@@ -0,0 +1,118 @@ ++#extension GL_EXT_control_flow_attributes : enable ++#extension GL_EXT_shader_16bit_storage : require ++#extension GL_EXT_shader_8bit_storage : require ++ ++#ifdef MUL_MAT_ID ++#define EXPERT_COUNT 8 ++#endif ++ ++#include "types.comp" ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; ++layout (binding = 1) readonly buffer BV2 {B_TYPE_VEC2 data_b_v2[];}; ++layout (binding = 1) readonly buffer BV4 {B_TYPE_VEC4 data_b_v4[];}; ++ ++layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; ++#ifdef MUL_MAT_ID ++layout (binding = 3) readonly buffer IDS {int data_ids[];}; ++#endif ++ ++#include "dequant_funcs.comp" ++ ++layout (push_constant) uniform parameter ++{ ++ uint ncols; ++ uint stride_a; ++ uint stride_b; ++ uint stride_d; ++ ++ uint batch_stride_a; ++ uint batch_stride_b; ++ uint batch_stride_d; ++ ++#ifdef MUL_MAT_ID ++ uint nei0; ++ uint ne11; ++#else ++ uint ne02; ++ uint ne12; ++ uint broadcast2; ++ uint broadcast3; ++#endif ++} p; ++ ++void get_offsets(out uint a_offset, out uint b_offset, out uint d_offset) { ++#ifdef MUL_MAT_ID ++ const uint expert_idx = gl_GlobalInvocationID.y; ++#else ++ const uint batch_idx = gl_GlobalInvocationID.y; ++#endif ++ ++#ifndef MUL_MAT_ID ++ uint batch_idx_a = 0; ++ if (batch_idx != 0) { ++ const uint i13 = batch_idx / p.ne12; ++ const uint i12 = batch_idx % p.ne12; ++ ++ const uint i03 = i13 / p.broadcast3; ++ const uint i02 = i12 / p.broadcast2; ++ ++ batch_idx_a = i03 * p.ne02 + i02; ++ } ++#else ++ const uint expert_id = data_ids[expert_idx]; ++#endif ++ ++ a_offset = ++#ifdef MUL_MAT_ID ++ expert_id * p.batch_stride_a; ++#else ++ batch_idx_a * p.batch_stride_a; ++#endif ++ b_offset = ++#ifdef MUL_MAT_ID ++ (expert_idx % p.ne11) * p.stride_b; ++#else ++ batch_idx * p.batch_stride_b; ++#endif ++ d_offset = ++#ifdef MUL_MAT_ID ++ expert_idx * p.stride_d; ++#else ++ batch_idx * p.batch_stride_d; ++#endif ++} ++ ++layout (constant_id = 0) const uint BLOCK_SIZE = 32; ++layout (constant_id = 1) const uint NUM_ROWS = 1; ++layout (constant_id = 2) const uint NUM_COLS = 1; ++ ++shared FLOAT_TYPE tmpsh[NUM_COLS][NUM_ROWS][BLOCK_SIZE]; ++ ++void reduce_result(const in FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const in uint32_t d_offset, const in uint32_t first_row, const in uint32_t num_rows, const in uint32_t tid) { ++ // sum up partial sums and write back result ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ [[unroll]] for (uint n = 0; n < num_rows; ++n) { ++ tmpsh[j][n][tid] = temp[j][n]; ++ } ++ } ++ barrier(); ++ [[unroll]] for (uint s = BLOCK_SIZE/2; s > 0; s >>= 1) { ++ if (tid < s) { ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ [[unroll]] for (uint n = 0; n < num_rows; ++n) { ++ tmpsh[j][n][tid] += tmpsh[j][n][tid + s]; ++ } ++ } ++ } ++ barrier(); ++ } ++ if (tid == 0) { ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ [[unroll]] for (uint n = 0; n < num_rows; ++n) { ++ data_d[j*p.batch_stride_d + d_offset + first_row + n] = D_TYPE(tmpsh[j][n][0]); ++ } ++ } ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp +new file mode 100644 +index 00000000..1cc4996d +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp +@@ -0,0 +1,71 @@ ++#version 450 ++ ++#extension GL_EXT_control_flow_attributes : enable ++#extension GL_EXT_shader_16bit_storage : require ++ ++#define BLOCK_SIZE 32 ++#define FLOAT_TYPE float ++ ++layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; ++layout (binding = 2) writeonly buffer D {D_TYPE dst[];}; ++ ++layout (push_constant) uniform parameter ++{ ++ uint ncols_x; ++ uint nrows_x; ++ uint row_stride_x; ++ uint channel_stride_x; ++ uint channel_x_divisor; ++ uint b_offset; ++ uint d_offset; ++} p; ++ ++shared FLOAT_TYPE tmp[BLOCK_SIZE]; ++ ++void main() { ++ const uint tid = gl_LocalInvocationID.x; ++ const uint row_x = gl_GlobalInvocationID.y; ++ const uint channel = gl_GlobalInvocationID.z; ++ const uint channel_x = channel / p.channel_x_divisor; ++ ++ const uint nrows_y = p.ncols_x; ++ const uint nrows_dst = p.nrows_x; ++ const uint row_dst = row_x; ++ ++ const uint idst = channel*nrows_dst + row_dst; ++ ++ tmp[tid] = 0.0f; ++ ++ for (uint col_x0 = 0; col_x0 < p.ncols_x; col_x0 += BLOCK_SIZE) { ++ const uint col_x = col_x0 + tid; ++ ++ if (col_x >= p.ncols_x) { ++ break; ++ } ++ ++ const uint row_y = col_x; ++ ++ const uint ix = channel_x*p.channel_stride_x + row_x*p.row_stride_x + col_x; ++ const uint iy = channel*nrows_y + row_y; ++ ++ const FLOAT_TYPE xi = FLOAT_TYPE(data_a[ix]); ++ ++ tmp[tid] = fma(xi, FLOAT_TYPE(data_b[iy]), tmp[tid]); ++ } ++ ++ // sum up partial sums and write back result ++ barrier(); ++ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { ++ if (tid < s) { ++ tmp[tid] += tmp[tid + s]; ++ } ++ barrier(); ++ } ++ ++ if (tid == 0) { ++ dst[idst] = tmp[0]; ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp +new file mode 100644 +index 00000000..9b443807 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp +@@ -0,0 +1,73 @@ ++#version 450 ++ ++#extension GL_EXT_control_flow_attributes : enable ++#extension GL_EXT_shader_16bit_storage : require ++ ++#define BLOCK_SIZE 32 ++#define FLOAT_TYPE float ++ ++layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; ++layout (binding = 2) writeonly buffer D {D_TYPE dst[];}; ++ ++layout (push_constant) uniform parameter ++{ ++ uint ncols_x; ++ uint nrows_x; ++ uint nchannels_x; ++ uint nchannels_y; ++ uint b_offset; ++ uint d_offset; ++} p; ++ ++shared FLOAT_TYPE tmp[BLOCK_SIZE]; ++ ++void main() { ++ const uint tid = gl_LocalInvocationID.x; ++ const uint row_x = gl_GlobalInvocationID.y; ++ const uint channel = gl_GlobalInvocationID.z; ++ const uint channel_x = channel / (p.nchannels_y / p.nchannels_x); ++ ++ const uint nrows_y = p.ncols_x; ++ const uint nrows_dst = p.nrows_x; ++ const uint row_dst = row_x; ++ ++ tmp[tid] = FLOAT_TYPE(0.0f); ++ ++ for (uint col_x0 = 0; col_x0 < p.ncols_x; col_x0 += BLOCK_SIZE) { ++ const uint col_x = col_x0 + tid; ++ ++ if (col_x >= p.ncols_x) { ++ break; ++ } ++ ++ // x is transposed and permuted ++ const uint ix = row_x*p.nchannels_x*p.ncols_x + channel_x*p.ncols_x + col_x; ++ const FLOAT_TYPE xi = FLOAT_TYPE(data_a[ix]); ++ ++ const uint row_y = col_x; ++ ++ // y is not transposed but permuted ++ const uint iy = channel*nrows_y + row_y; ++ ++ tmp[tid] = fma(xi, FLOAT_TYPE(data_b[iy]), tmp[tid]); ++ } ++ ++ // dst is not transposed and not permuted ++ const uint idst = channel*nrows_dst + row_dst; ++ ++ // sum up partial sums and write back result ++ barrier(); ++ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { ++ if (tid < s) { ++ tmp[tid] += tmp[tid + s]; ++ } ++ barrier(); ++ } ++ ++ if (tid == 0) { ++ dst[idst] = tmp[0]; ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp +new file mode 100644 +index 00000000..93421344 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp +@@ -0,0 +1,115 @@ ++#version 450 ++#extension GL_EXT_shader_explicit_arithmetic_types : require ++ ++#include "mul_mat_vec_base.comp" ++ ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { ++ uint a_offset, b_offset, d_offset; ++ get_offsets(a_offset, b_offset, d_offset); ++ ++ const uint num_blocks_per_row = p.ncols / QUANT_K; ++ ++ // 16 threads are used to process each block ++ const uint it_size = gl_WorkGroupSize.x/16; ++ const uint tid = gl_LocalInvocationID.x; ++ const uint itid = tid%16; // 0...16 ++ const uint ix = tid/16; ++ ++ const uint step = 8; ++ ++ const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... ++ const uint v_in = itid - step*v_im; // 0...15 or 0...7 ++ ++ const uint l0 = 2*v_in; // 0...15 ++ const uint q_offset = 32*v_im + l0; ++ const uint s_offset = 8*v_im; ++ const uint y_offset = 128*v_im + l0; ++ ++ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; ++ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { ++ temp[j][i] = FLOAT_TYPE(0); ++ } ++ } ++ ++ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { ++ const uint y_idx = i * QUANT_K + y_offset; ++ ++ [[unroll]] for (uint n = 0; n < num_rows; ++n) { ++ const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; ++ f16vec2 d = data_a[ib0 + i].d; ++ const FLOAT_TYPE dall = d.x; ++ const FLOAT_TYPE dmin = d.y; ++ ++ uint32_t s0_u32 = data_a_packed32[ib0 + i].scales[s_offset / 4 + 0]; ++ uint32_t s4_u32 = data_a_packed32[ib0 + i].scales[s_offset / 4 + 1]; ++ ++ uint32_t s0_lo4_u32 = s0_u32 & 0x0F0F0F0F; ++ uint32_t s0_hi4_u32 = (s0_u32 >> 4) & 0x0F0F0F0F; ++ uint32_t s4_lo4_u32 = s4_u32 & 0x0F0F0F0F; ++ uint32_t s4_hi4_u32 = (s4_u32 >> 4) & 0x0F0F0F0F; ++ ++ uvec4 s0_lo4 = uvec4(unpack8(s0_lo4_u32)); ++ uvec4 s4_lo4 = uvec4(unpack8(s4_lo4_u32)); ++ uvec4 s0_hi4 = uvec4(unpack8(s0_hi4_u32)); ++ uvec4 s4_hi4 = uvec4(unpack8(s4_hi4_u32)); ++ ++ uint16_t qs0_u16 = data_a_packed16[ib0 + i].qs[q_offset / 2 + 0]; ++ uint16_t qs16_u16 = data_a_packed16[ib0 + i].qs[q_offset / 2 + 8]; ++ uvec2 qs0 = uvec2(unpack8(qs0_u16)); ++ uvec2 qs16 = uvec2(unpack8(qs16_u16)); ++ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ B_TYPE_VEC2 b0 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 0]; ++ B_TYPE_VEC2 b16 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 8]; ++ B_TYPE_VEC2 b32 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 16]; ++ B_TYPE_VEC2 b48 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 24]; ++ B_TYPE_VEC2 b64 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 32]; ++ B_TYPE_VEC2 b80 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 40]; ++ B_TYPE_VEC2 b96 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 48]; ++ B_TYPE_VEC2 b112 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 56]; ++ ++ FLOAT_TYPE sum1 = FLOAT_TYPE(0.0); ++ FLOAT_TYPE sum2 = FLOAT_TYPE(0.0); ++ [[unroll]] for (int l = 0; l < 2; ++l) { ++ sum1 = fma(FLOAT_TYPE(b0[l]), FLOAT_TYPE(s0_lo4[0]) * FLOAT_TYPE((qs0[l] >> 0) & 3), ++ fma(FLOAT_TYPE(b16[l]), FLOAT_TYPE(s0_lo4[1]) * FLOAT_TYPE((qs16[l] >> 0) & 3), ++ fma(FLOAT_TYPE(b32[l]), FLOAT_TYPE(s0_lo4[2]) * FLOAT_TYPE((qs0[l] >> 2) & 3), ++ fma(FLOAT_TYPE(b48[l]), FLOAT_TYPE(s0_lo4[3]) * FLOAT_TYPE((qs16[l] >> 2) & 3), ++ fma(FLOAT_TYPE(b64[l]), FLOAT_TYPE(s4_lo4[0]) * FLOAT_TYPE((qs0[l] >> 4) & 3), ++ fma(FLOAT_TYPE(b80[l]), FLOAT_TYPE(s4_lo4[1]) * FLOAT_TYPE((qs16[l] >> 4) & 3), ++ fma(FLOAT_TYPE(b96[l]), FLOAT_TYPE(s4_lo4[2]) * FLOAT_TYPE((qs0[l] >> 6) & 3), ++ fma(FLOAT_TYPE(b112[l]), FLOAT_TYPE(s4_lo4[3]) * FLOAT_TYPE((qs16[l] >> 6) & 3), sum1)))))))); ++ sum2 = fma(FLOAT_TYPE(b0[l]), FLOAT_TYPE(s0_hi4[0]), ++ fma(FLOAT_TYPE(b16[l]), FLOAT_TYPE(s0_hi4[1]), ++ fma(FLOAT_TYPE(b32[l]), FLOAT_TYPE(s0_hi4[2]), ++ fma(FLOAT_TYPE(b48[l]), FLOAT_TYPE(s0_hi4[3]), ++ fma(FLOAT_TYPE(b64[l]), FLOAT_TYPE(s4_hi4[0]), ++ fma(FLOAT_TYPE(b80[l]), FLOAT_TYPE(s4_hi4[1]), ++ fma(FLOAT_TYPE(b96[l]), FLOAT_TYPE(s4_hi4[2]), ++ fma(FLOAT_TYPE(b112[l]), FLOAT_TYPE(s4_hi4[3]), sum2)))))))); ++ } ++ temp[j][n] = fma(dall, sum1, fma(-dmin, sum2, temp[j][n])); ++ } ++ } ++ } ++ ++ reduce_result(temp, d_offset, first_row, num_rows, tid); ++} ++ ++void main() { ++ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); ++ ++ // do NUM_ROWS at a time, unless there aren't enough remaining rows ++ if (first_row + NUM_ROWS <= p.stride_d) { ++ compute_outputs(first_row, NUM_ROWS); ++ } else { ++ if (first_row >= p.stride_d) { ++ return; ++ } ++ compute_outputs(first_row, p.stride_d - first_row); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp +new file mode 100644 +index 00000000..86b0159d +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp +@@ -0,0 +1,103 @@ ++#version 450 ++#extension GL_EXT_shader_explicit_arithmetic_types : require ++ ++#include "mul_mat_vec_base.comp" ++ ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { ++ uint a_offset, b_offset, d_offset; ++ get_offsets(a_offset, b_offset, d_offset); ++ ++ const uint num_blocks_per_row = p.ncols / QUANT_K; ++ ++ // 16 threads are used to process each block ++ const uint it_size = gl_WorkGroupSize.x/16; ++ const uint tid = gl_LocalInvocationID.x; ++ const uint itid = tid%16; // 0...16 ++ const uint ix = tid/16; ++ ++ const uint step = 8; ++ ++ const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... ++ const uint v_in = itid - step*v_im; // 0...15 or 0...7 ++ ++ const uint8_t m = uint8_t(1 << (4 * v_im)); ++ ++ const uint l0 = 2*v_in; // 0...15 ++ const uint q_offset = 32*v_im + l0; ++ const uint y_offset = 128*v_im + l0; ++ ++ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; ++ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { ++ temp[j][i] = FLOAT_TYPE(0); ++ } ++ } ++ ++ const uint s_shift = 4 * v_im; ++ ++ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { ++ const uint y_idx = i * QUANT_K + y_offset; ++ ++ [[unroll]] for (uint n = 0; n < num_rows; ++n) { ++ const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; ++ const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d); ++ ++ uint16_t s0_16 = data_a_packed16[ib0 + i].scales[0]; ++ uint16_t s2_16 = data_a_packed16[ib0 + i].scales[1]; ++ uint16_t s4_16 = data_a_packed16[ib0 + i].scales[2]; ++ uint16_t s6_16 = data_a_packed16[ib0 + i].scales[3]; ++ uint16_t s8_16 = data_a_packed16[ib0 + i].scales[4]; ++ uint16_t s10_16 = data_a_packed16[ib0 + i].scales[5]; ++ u8vec2 s0 = unpack8(s0_16); ++ u8vec2 s2 = unpack8(s2_16); ++ u8vec2 s4 = unpack8(s4_16); ++ u8vec2 s6 = unpack8(s6_16); ++ u8vec2 s8 = unpack8(s8_16); ++ u8vec2 s10 = unpack8(s10_16); ++ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ ++ B_TYPE_VEC2 b0 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 0]; ++ B_TYPE_VEC2 b16 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 8]; ++ B_TYPE_VEC2 b32 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 16]; ++ B_TYPE_VEC2 b48 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 24]; ++ B_TYPE_VEC2 b64 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 32]; ++ B_TYPE_VEC2 b80 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 40]; ++ B_TYPE_VEC2 b96 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 48]; ++ B_TYPE_VEC2 b112 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 56]; ++ ++ FLOAT_TYPE sum = FLOAT_TYPE(0.0); ++ [[unroll]] for (int l = 0; l < 2; ++l) { ++ sum = fma(FLOAT_TYPE(b0[l]) * FLOAT_TYPE(int8_t(((s0[0] >> s_shift) & 0xF) | ((s8[0] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 0)) != 0) ? 0 : 4)), ++ fma(FLOAT_TYPE(b32[l]) * FLOAT_TYPE(int8_t(((s2[0] >> s_shift) & 0xF) | ((s10[0] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 1)) != 0) ? 0 : 4)), ++ fma(FLOAT_TYPE(b64[l]) * FLOAT_TYPE(int8_t(((s4[0] >> s_shift) & 0xF) | ((s8[0] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 2)) != 0) ? 0 : 4)), ++ fma(FLOAT_TYPE(b96[l]) * FLOAT_TYPE(int8_t(((s6[0] >> s_shift) & 0xF) | ((s10[0] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 3)) != 0) ? 0 : 4)), ++ fma(FLOAT_TYPE(b16[l]) * FLOAT_TYPE(int8_t(((s0[1] >> s_shift) & 0xF) | ((s8[1] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 0)) != 0) ? 0 : 4)), ++ fma(FLOAT_TYPE(b48[l]) * FLOAT_TYPE(int8_t(((s2[1] >> s_shift) & 0xF) | ((s10[1] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 1)) != 0) ? 0 : 4)), ++ fma(FLOAT_TYPE(b80[l]) * FLOAT_TYPE(int8_t(((s4[1] >> s_shift) & 0xF) | ((s8[1] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 2)) != 0) ? 0 : 4)), ++ fma(FLOAT_TYPE(b112[l]) * FLOAT_TYPE(int8_t(((s6[1] >> s_shift) & 0xF) | ((s10[1] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 3)) != 0) ? 0 : 4)), sum)))))))); ++ } ++ temp[j][n] = fma(d, sum, temp[j][n]); ++ } ++ } ++ } ++ ++ reduce_result(temp, d_offset, first_row, num_rows, tid); ++} ++ ++void main() { ++ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); ++ ++ // do NUM_ROWS at a time, unless there aren't enough remaining rows ++ if (first_row + NUM_ROWS <= p.stride_d) { ++ compute_outputs(first_row, NUM_ROWS); ++ } else { ++ if (first_row >= p.stride_d) { ++ return; ++ } ++ compute_outputs(first_row, p.stride_d - first_row); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp +new file mode 100644 +index 00000000..cd1dd8e8 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp +@@ -0,0 +1,133 @@ ++#version 450 ++ ++#extension GL_EXT_shader_explicit_arithmetic_types : require ++ ++#include "mul_mat_vec_base.comp" ++ ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { ++ uint a_offset, b_offset, d_offset; ++ get_offsets(a_offset, b_offset, d_offset); ++ ++ const uint num_blocks_per_row = p.ncols / QUANT_K; ++ ++ // 16 threads are used to process each block ++ const uint it_size = gl_WorkGroupSize.x/16; ++ const uint tid = gl_LocalInvocationID.x; ++ const uint itid = tid%16; // 0...16 ++ const uint ix = tid/16; ++ ++ const uint step = 4; ++ ++ const uint il = itid/step; // 0...3 ++ const uint ir = itid - step*il; // 0...7 or 0...3 ++ const uint n = 4; ++ ++ const uint v_im = il / 2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224 ++ const uint v_in = il % 2; ++ ++ const uint l0 = n * (2 * ir + v_in); // 0...15 ++ const uint q_offset = 32*v_im + l0; ++ const uint y_offset = 64*v_im + l0; ++ ++ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; ++ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { ++ temp[j][i] = FLOAT_TYPE(0); ++ } ++ } ++ ++ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { ++ const uint y1_idx = i * QUANT_K + y_offset; ++ const uint y2_idx = y1_idx + 128; ++ ++ [[unroll]] for (uint n = 0; n < num_rows; ++n) { ++ const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; ++ f16vec2 d = data_a[ib0 + i].d; ++ const FLOAT_TYPE dall = FLOAT_TYPE(d.x); ++ const FLOAT_TYPE dmin = FLOAT_TYPE(d.y); ++ ++ uint32_t scale0_u32 = data_a_packed16[ib0 + i].scales[v_im ]; ++ uint32_t scale4_u32 = data_a_packed16[ib0 + i].scales[v_im + 2]; ++ uint32_t scale8_u32 = data_a_packed16[ib0 + i].scales[v_im + 4]; ++ uvec4 scale0 = uvec4(unpack8(scale0_u32)); ++ uvec4 scale4 = uvec4(unpack8(scale4_u32)); ++ uvec4 scale8 = uvec4(unpack8(scale8_u32)); ++ ++ const uint32_t sc0 = ( scale0.x & 0x3f); ++ const uint32_t sc1 = ( scale0.y & 0x3f); ++ const uint32_t sc2 = ( scale4.x & 0x3f); ++ const uint32_t sc3 = ( scale4.y & 0x3f); ++ const uint32_t sc4 = (( scale8.x & 0x0f) | ((scale0.x & 0xc0) >> 2)); ++ const uint32_t sc5 = (( scale8.y & 0x0f) | ((scale0.y & 0xc0) >> 2)); ++ const uint32_t sc6 = (((scale8.x >> 4) & 0x0f) | ((scale4.x & 0xc0) >> 2)); ++ const uint32_t sc7 = (((scale8.y >> 4) & 0x0f) | ((scale4.y & 0xc0) >> 2)); ++ ++ uint32_t qs0_u32 = data_a_packed32[ib0 + i].qs[q_offset / 4]; ++ uint32_t qs64_u32 = data_a_packed32[ib0 + i].qs[q_offset / 4 + 16]; ++ ++ uint32_t qs0_u32_lo4 = qs0_u32 & 0x0F0F0F0F; ++ uint32_t qs0_u32_hi4 = (qs0_u32 >> 4) & 0x0F0F0F0F; ++ uint32_t qs64_u32_lo4 = qs64_u32 & 0x0F0F0F0F; ++ uint32_t qs64_u32_hi4 = (qs64_u32 >> 4) & 0x0F0F0F0F; ++ ++ uvec4 qs0_lo4 = uvec4(unpack8(qs0_u32_lo4)); ++ uvec4 qs64_lo4 = uvec4(unpack8(qs64_u32_lo4)); ++ uvec4 qs0_hi4 = uvec4(unpack8(qs0_u32_hi4)); ++ uvec4 qs64_hi4 = uvec4(unpack8(qs64_u32_hi4)); ++ ++ const uint32_t q4_0 = qs0_lo4.x; ++ const uint32_t q4_1 = qs0_lo4.y; ++ const uint32_t q4_2 = qs0_lo4.z; ++ const uint32_t q4_3 = qs0_lo4.w; ++ const uint32_t q4_4 = qs0_hi4.x; ++ const uint32_t q4_5 = qs0_hi4.y; ++ const uint32_t q4_6 = qs0_hi4.z; ++ const uint32_t q4_7 = qs0_hi4.w; ++ const uint32_t q4_8 = qs64_lo4.x; ++ const uint32_t q4_9 = qs64_lo4.y; ++ const uint32_t q4_10 = qs64_lo4.z; ++ const uint32_t q4_11 = qs64_lo4.w; ++ const uint32_t q4_12 = qs64_hi4.x; ++ const uint32_t q4_13 = qs64_hi4.y; ++ const uint32_t q4_14 = qs64_hi4.z; ++ const uint32_t q4_15 = qs64_hi4.w; ++ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ B_TYPE_VEC4 by10 = data_b_v4[(j*p.batch_stride_b + b_offset + y1_idx) / 4]; ++ B_TYPE_VEC4 by132 = data_b_v4[(j*p.batch_stride_b + b_offset + y1_idx) / 4 + 8]; ++ B_TYPE_VEC4 by20 = data_b_v4[(j*p.batch_stride_b + b_offset + y2_idx) / 4]; ++ B_TYPE_VEC4 by232 = data_b_v4[(j*p.batch_stride_b + b_offset + y2_idx) / 4 + 8]; ++ ++ const FLOAT_TYPE sx = fma(FLOAT_TYPE(by10.x), q4_0, fma(FLOAT_TYPE(by10.y), q4_1, fma(FLOAT_TYPE(by10.z), q4_2, FLOAT_TYPE(by10.w) * q4_3))); ++ const FLOAT_TYPE sy = fma(FLOAT_TYPE(by132.x), q4_4, fma(FLOAT_TYPE(by132.y), q4_5, fma(FLOAT_TYPE(by132.z), q4_6, FLOAT_TYPE(by132.w) * q4_7))); ++ const FLOAT_TYPE sz = fma(FLOAT_TYPE(by20.x), q4_8, fma(FLOAT_TYPE(by20.y), q4_9, fma(FLOAT_TYPE(by20.z), q4_10, FLOAT_TYPE(by20.w) * q4_11))); ++ const FLOAT_TYPE sw = fma(FLOAT_TYPE(by232.x), q4_12, fma(FLOAT_TYPE(by232.y), q4_13, fma(FLOAT_TYPE(by232.z), q4_14, FLOAT_TYPE(by232.w) * q4_15))); ++ const FLOAT_TYPE smin = ++ fma(FLOAT_TYPE(by10.x), sc2, fma(FLOAT_TYPE(by132.x), sc3, fma(FLOAT_TYPE(by20.x), sc6, fma(FLOAT_TYPE(by232.x), sc7, ++ fma(FLOAT_TYPE(by10.y), sc2, fma(FLOAT_TYPE(by132.y), sc3, fma(FLOAT_TYPE(by20.y), sc6, fma(FLOAT_TYPE(by232.y), sc7, ++ fma(FLOAT_TYPE(by10.z), sc2, fma(FLOAT_TYPE(by132.z), sc3, fma(FLOAT_TYPE(by20.z), sc6, fma(FLOAT_TYPE(by232.z), sc7, ++ fma(FLOAT_TYPE(by10.w), sc2, fma(FLOAT_TYPE(by132.w), sc3, fma(FLOAT_TYPE(by20.w), sc6, FLOAT_TYPE(by232.w) * sc7))))))))))))))); ++ temp[j][n] = fma(dall, fma(sx, sc0, fma(sy, sc1, fma(sz, sc4, sw * sc5))), fma(-dmin, smin, temp[j][n])); ++ } ++ } ++ } ++ ++ reduce_result(temp, d_offset, first_row, num_rows, tid); ++} ++ ++void main() { ++ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); ++ ++ // do NUM_ROWS at a time, unless there aren't enough remaining rows ++ if (first_row + NUM_ROWS <= p.stride_d) { ++ compute_outputs(first_row, NUM_ROWS); ++ } else { ++ if (first_row >= p.stride_d) { ++ return; ++ } ++ compute_outputs(first_row, p.stride_d - first_row); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp +new file mode 100644 +index 00000000..0a68891c +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp +@@ -0,0 +1,162 @@ ++#version 450 ++ ++#extension GL_EXT_shader_explicit_arithmetic_types : require ++ ++#include "mul_mat_vec_base.comp" ++ ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { ++ uint a_offset, b_offset, d_offset; ++ get_offsets(a_offset, b_offset, d_offset); ++ ++ const uint num_blocks_per_row = p.ncols / QUANT_K; ++ ++ // 16 threads are used to process each block ++ const uint it_size = gl_WorkGroupSize.x/16; ++ const uint tid = gl_LocalInvocationID.x; ++ const uint itid = tid%16; // 0...16 ++ const uint ix = tid/16; ++ ++ const uint il = itid/4; // 0...3 ++ const uint ir = itid - 4*il; // 0...7 or 0...3 ++ ++ const uint v_im = il / 2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224 ++ const uint v_in = il % 2; ++ ++ const uint l0 = 4*ir + 2*v_in; // 0...15 ++ const uint q_offset = 32*v_im + l0; ++ const uint y_offset = 64*v_im + l0; ++ ++ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; ++ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { ++ temp[j][i] = FLOAT_TYPE(0); ++ } ++ } ++ ++ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { ++ const uint y1_idx = i * QUANT_K + y_offset; ++ const uint y2_idx = y1_idx + 128; ++ ++ [[unroll]] for (uint n = 0; n < num_rows; ++n) { ++ const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; ++ f16vec2 d = data_a[ib0 + i].d; ++ const FLOAT_TYPE dall = FLOAT_TYPE(d.x); ++ const FLOAT_TYPE dmin = FLOAT_TYPE(d.y); ++ ++ uint32_t scale0_u32 = data_a_packed16[ib0 + i].scales[v_im ]; ++ uint32_t scale4_u32 = data_a_packed16[ib0 + i].scales[v_im + 2]; ++ uint32_t scale8_u32 = data_a_packed16[ib0 + i].scales[v_im + 4]; ++ uvec4 scale0 = uvec4(unpack8(scale0_u32)); ++ uvec4 scale4 = uvec4(unpack8(scale4_u32)); ++ uvec4 scale8 = uvec4(unpack8(scale8_u32)); ++ ++ const uint32_t sc0 = ( scale0.x & 0x3f); ++ const uint32_t sc1 = ( scale0.y & 0x3f); ++ const uint32_t sc2 = ( scale4.x & 0x3f); ++ const uint32_t sc3 = ( scale4.y & 0x3f); ++ const uint32_t sc4 = (( scale8.x & 0x0f) | ((scale0.x & 0xc0) >> 2)); ++ const uint32_t sc5 = (( scale8.y & 0x0f) | ((scale0.y & 0xc0) >> 2)); ++ const uint32_t sc6 = (((scale8.x >> 4) & 0x0f) | ((scale4.x & 0xc0) >> 2)); ++ const uint32_t sc7 = (((scale8.y >> 4) & 0x0f) | ((scale4.y & 0xc0) >> 2)); ++ ++ uint32_t qs0_16_u32 = uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 8]) << 16); ++ uint32_t qs64_80_u32 = uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 32]) | (uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 40]) << 16); ++ ++ uint32_t qs0_16_u32_lo4 = qs0_16_u32 & 0x0F0F0F0F; ++ uint32_t qs0_16_u32_hi4 = (qs0_16_u32 >> 4) & 0x0F0F0F0F; ++ uint32_t qs64_80_u32_lo4 = qs64_80_u32 & 0x0F0F0F0F; ++ uint32_t qs64_80_u32_hi4 = (qs64_80_u32 >> 4) & 0x0F0F0F0F; ++ ++ uint32_t qh = pack32(u16vec2(data_a_packed16[ib0 + i].qh[l0 / 2], data_a_packed16[ib0 + i].qh[l0 / 2 + 8])); ++ ++ uint32_t qs0_16_lo4_offset16 = ((qh >> (2*v_im)) & 0x01010101) << 4; ++ uint32_t qs0_16_hi4_offset16 = ((qh >> (2*v_im)) & 0x02020202) << 3; ++ uint32_t qs64_80_lo4_offset16 = ((qh >> (2*v_im)) & 0x10101010) << 0; ++ uint32_t qs64_80_hi4_offset16 = ((qh >> (2*v_im)) & 0x20202020) >> 1; ++ ++ qs0_16_u32_lo4 += qs0_16_lo4_offset16; ++ qs0_16_u32_hi4 += qs0_16_hi4_offset16; ++ qs64_80_u32_lo4 += qs64_80_lo4_offset16; ++ qs64_80_u32_hi4 += qs64_80_hi4_offset16; ++ ++ uvec4 qs0_16_lo4 = uvec4(unpack8(qs0_16_u32_lo4)); ++ uvec4 qs64_80_lo4 = uvec4(unpack8(qs64_80_u32_lo4)); ++ uvec4 qs0_16_hi4 = uvec4(unpack8(qs0_16_u32_hi4)); ++ uvec4 qs64_80_hi4 = uvec4(unpack8(qs64_80_u32_hi4)); ++ ++ const uint32_t q4_0 = qs0_16_lo4.x; ++ const uint32_t q4_1 = qs0_16_lo4.y; ++ const uint32_t q4_2 = qs0_16_lo4.z; ++ const uint32_t q4_3 = qs0_16_lo4.w; ++ const uint32_t q4_4 = qs0_16_hi4.x; ++ const uint32_t q4_5 = qs0_16_hi4.y; ++ const uint32_t q4_6 = qs0_16_hi4.z; ++ const uint32_t q4_7 = qs0_16_hi4.w; ++ const uint32_t q4_8 = qs64_80_lo4.x; ++ const uint32_t q4_9 = qs64_80_lo4.y; ++ const uint32_t q4_10 = qs64_80_lo4.z; ++ const uint32_t q4_11 = qs64_80_lo4.w; ++ const uint32_t q4_12 = qs64_80_hi4.x; ++ const uint32_t q4_13 = qs64_80_hi4.y; ++ const uint32_t q4_14 = qs64_80_hi4.z; ++ const uint32_t q4_15 = qs64_80_hi4.w; ++ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ B_TYPE_VEC2 by10 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2]; ++ B_TYPE_VEC2 by116 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 8]; ++ B_TYPE_VEC2 by132 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 16]; ++ B_TYPE_VEC2 by148 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 24]; ++ B_TYPE_VEC2 by20 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2]; ++ B_TYPE_VEC2 by216 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 8]; ++ B_TYPE_VEC2 by232 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 16]; ++ B_TYPE_VEC2 by248 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 24]; ++ ++ const FLOAT_TYPE sx = ++ fma(FLOAT_TYPE(by10.x), q4_0, ++ fma(FLOAT_TYPE(by10.y), q4_1, ++ fma(FLOAT_TYPE(by116.x), q4_2, ++ FLOAT_TYPE(by116.y) * q4_3))); ++ const FLOAT_TYPE sy = ++ fma(FLOAT_TYPE(by132.x), q4_4, ++ fma(FLOAT_TYPE(by132.y), q4_5, ++ fma(FLOAT_TYPE(by148.x), q4_6, ++ FLOAT_TYPE(by148.y) * q4_7))); ++ const FLOAT_TYPE sz = ++ fma(FLOAT_TYPE(by20.x), q4_8, ++ fma(FLOAT_TYPE(by20.y), q4_9, ++ fma(FLOAT_TYPE(by216.x), q4_10, ++ FLOAT_TYPE(by216.y) * q4_11))); ++ const FLOAT_TYPE sw = ++ fma(FLOAT_TYPE(by232.x), q4_12, ++ fma(FLOAT_TYPE(by232.y), q4_13, ++ fma(FLOAT_TYPE(by248.x), q4_14, ++ FLOAT_TYPE(by248.y) * q4_15))); ++ const FLOAT_TYPE smin = ++ fma(FLOAT_TYPE(by10.x) + FLOAT_TYPE(by10.y) + FLOAT_TYPE(by116.x) + FLOAT_TYPE(by116.y), sc2, ++ fma(FLOAT_TYPE(by132.x) + FLOAT_TYPE(by132.y) + FLOAT_TYPE(by148.x) + FLOAT_TYPE(by148.y), sc3, ++ fma(FLOAT_TYPE(by20.x) + FLOAT_TYPE(by20.y) + FLOAT_TYPE(by216.x) + FLOAT_TYPE(by216.y), sc6, ++ (FLOAT_TYPE(by232.x) + FLOAT_TYPE(by232.y) + FLOAT_TYPE(by248.x) + FLOAT_TYPE(by248.y)) * sc7))); ++ temp[j][n] = fma(dall, fma(sx, sc0, fma(sy, sc1, fma(sz, sc4, sw * sc5))), fma(-dmin, smin, temp[j][n])); ++ } ++ } ++ } ++ ++ reduce_result(temp, d_offset, first_row, num_rows, tid); ++} ++ ++void main() { ++ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); ++ ++ // do NUM_ROWS at a time, unless there aren't enough remaining rows ++ if (first_row + NUM_ROWS <= p.stride_d) { ++ compute_outputs(first_row, NUM_ROWS); ++ } else { ++ if (first_row >= p.stride_d) { ++ return; ++ } ++ compute_outputs(first_row, p.stride_d - first_row); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp +new file mode 100644 +index 00000000..70e13a56 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp +@@ -0,0 +1,112 @@ ++#version 450 ++ ++#extension GL_EXT_shader_explicit_arithmetic_types : require ++ ++#include "mul_mat_vec_base.comp" ++ ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { ++ uint a_offset, b_offset, d_offset; ++ get_offsets(a_offset, b_offset, d_offset); ++ ++ const uint num_blocks_per_row = p.ncols / QUANT_K; ++ ++ // 16 threads are used to process each block ++ const uint it_size = gl_WorkGroupSize.x/16; ++ const uint tid = gl_LocalInvocationID.x; ++ const uint itid = tid%16; // 0...16 ++ const uint ix = tid/16; ++ ++ const uint step = 8; ++ ++ const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... ++ const uint v_in = itid - step*v_im; // 0...15 or 0...7 ++ ++ const uint l0 = 4 * v_in; // 0, 4, 8, ..., 28 ++ const uint is = v_in / 4; ++ ++ const uint ql_offset = 64*v_im + l0; ++ const uint qh_offset = 32*v_im + l0; ++ const uint s_offset = 8*v_im + is; ++ const uint y_offset = 128*v_im + l0; ++ ++ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; ++ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { ++ temp[j][i] = FLOAT_TYPE(0); ++ } ++ } ++ ++ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { ++ const uint y_idx = i * QUANT_K + y_offset; ++ ++ [[unroll]] for (uint n = 0; n < num_rows; ++n) { ++ const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; ++ const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d); ++ ++ FLOAT_TYPE scales[4]; ++ scales[0] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 0]); ++ scales[1] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 2]); ++ scales[2] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 4]); ++ scales[3] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 6]); ++ ++ uint32_t ql0_u32 = uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 1]) << 16); ++ uint32_t ql32_u32 = uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 16]) | (uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 17]) << 16); ++ ++ uint32_t ql0_u32_lo4 = ql0_u32 & 0x0F0F0F0F; ++ uint32_t ql0_u32_hi4 = (ql0_u32 >> 4) & 0x0F0F0F0F; ++ uint32_t ql32_u32_lo4 = ql32_u32 & 0x0F0F0F0F; ++ uint32_t ql32_u32_hi4 = (ql32_u32 >> 4) & 0x0F0F0F0F; ++ ++ uint32_t qh_u32 = uint32_t(data_a_packed16[ib0 + i].qh[qh_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].qh[qh_offset / 2 + 1]) << 16); ++ uint32_t qh0_u32 = (qh_u32 & 0x03030303) << 4; ++ uint32_t qh2_u32 = (qh_u32 & 0x0C0C0C0C) << 2; ++ uint32_t qh4_u32 = (qh_u32 & 0x30303030) << 0; ++ uint32_t qh6_u32 = (qh_u32 & 0xC0C0C0C0) >> 2; ++ ++ uint32_t q0_u32 = ql0_u32_lo4 | qh0_u32; ++ uint32_t q1_u32 = ql32_u32_lo4 | qh2_u32; ++ uint32_t q2_u32 = ql0_u32_hi4 | qh4_u32; ++ uint32_t q3_u32 = ql32_u32_hi4 | qh6_u32; ++ ++ uvec4 q0 = uvec4(unpack8(q0_u32)); ++ uvec4 q1 = uvec4(unpack8(q1_u32)); ++ uvec4 q2 = uvec4(unpack8(q2_u32)); ++ uvec4 q3 = uvec4(unpack8(q3_u32)); ++ ++ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { ++ B_TYPE_VEC4 by0 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4]; ++ B_TYPE_VEC4 by32 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 8]; ++ B_TYPE_VEC4 by64 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 16]; ++ B_TYPE_VEC4 by96 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 24]; ++ ++ FLOAT_TYPE sum = FLOAT_TYPE(0.0); ++ [[unroll]] for (int l = 0; l < 4; ++l) { ++ sum = fma(FLOAT_TYPE(by0[l]) * scales[0], FLOAT_TYPE(int8_t(q0[l]) - 32), ++ fma(FLOAT_TYPE(by32[l]) * scales[1], FLOAT_TYPE(int8_t(q1[l]) - 32), ++ fma(FLOAT_TYPE(by64[l]) * scales[2], FLOAT_TYPE(int8_t(q2[l]) - 32), ++ fma(FLOAT_TYPE(by96[l]) * scales[3], FLOAT_TYPE(int8_t(q3[l]) - 32), sum)))); ++ } ++ temp[j][n] += sum * d; ++ } ++ } ++ } ++ ++ reduce_result(temp, d_offset, first_row, num_rows, tid); ++} ++ ++void main() { ++ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); ++ ++ // do NUM_ROWS at a time, unless there aren't enough remaining rows ++ if (first_row + NUM_ROWS <= p.stride_d) { ++ compute_outputs(first_row, NUM_ROWS); ++ } else { ++ if (first_row >= p.stride_d) { ++ return; ++ } ++ compute_outputs(first_row, p.stride_d - first_row); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +new file mode 100644 +index 00000000..48122cbe +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +@@ -0,0 +1,631 @@ ++#version 450 ++ ++#extension GL_EXT_control_flow_attributes : enable ++#extension GL_EXT_shader_16bit_storage : require ++ ++#ifdef FLOAT16 ++#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require ++#endif ++ ++#ifdef COOPMAT ++#extension GL_KHR_cooperative_matrix : enable ++#extension GL_KHR_memory_scope_semantics : enable ++#extension GL_KHR_shader_subgroup_basic : enable ++#endif ++ ++#ifdef MUL_MAT_ID ++#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require ++#endif ++ ++#include "types.comp" ++ ++#ifndef LOAD_VEC_A ++#define LOAD_VEC_A 1 ++#endif ++#ifndef LOAD_VEC_B ++#define LOAD_VEC_B 1 ++#endif ++ ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; ++layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; ++ ++#ifdef MUL_MAT_ID ++layout (binding = 3) readonly buffer IDS {int data_ids[];}; ++#endif ++ ++layout (push_constant) uniform parameter ++{ ++ uint M; ++ uint N; ++ uint K; ++ uint stride_a; ++ uint stride_b; ++ uint stride_d; ++ ++ uint batch_stride_a; ++ uint batch_stride_b; ++ uint batch_stride_d; ++ ++#ifdef MUL_MAT_ID ++ uint nei0; ++ uint nei1; ++ uint nbi1; ++ uint ne11; ++#else ++ uint k_split; ++ uint ne02; ++ uint ne12; ++ uint broadcast2; ++ uint broadcast3; ++#endif ++} p; ++ ++layout (constant_id = 0) const uint BLOCK_SIZE = 64; ++layout (constant_id = 1) const uint BM = 64; ++layout (constant_id = 2) const uint BN = 64; ++layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working with a quant ++layout (constant_id = 4) const uint WM = 32; ++layout (constant_id = 5) const uint WN = 32; ++layout (constant_id = 6) const uint WMITER = 2; ++layout (constant_id = 7) const uint TM = 4; ++layout (constant_id = 8) const uint TN = 2; ++layout (constant_id = 9) const uint TK = 1; // Only needed for coopmat ++layout (constant_id = 10) const uint WARP = 32; ++ ++#ifdef COOPMAT ++#define SHMEM_STRIDE (BK + 8) ++#else ++#define SHMEM_STRIDE (BK + 1) ++#endif ++ ++shared FLOAT_TYPE buf_a[BM * SHMEM_STRIDE]; ++shared FLOAT_TYPE buf_b[BN * SHMEM_STRIDE]; ++ ++#ifdef MUL_MAT_ID ++shared u16vec2 row_ids[3072]; ++#endif // MUL_MAT_ID ++ ++#define NUM_WARPS (BLOCK_SIZE / WARP) ++ ++#ifdef COOPMAT ++shared ACC_TYPE coopmat_stage[TM * TN * NUM_WARPS]; ++#endif ++ ++void main() { ++#if defined(DATA_A_IQ4_NL) ++ init_iq4nl_shmem(); ++#endif ++ ++#ifdef MUL_MAT_ID ++ const uint expert_idx = gl_GlobalInvocationID.z; ++#else ++ const uint batch_idx = gl_GlobalInvocationID.z; ++ ++ const uint i13 = batch_idx / p.ne12; ++ const uint i12 = batch_idx % p.ne12; ++ ++ const uint i03 = i13 / p.broadcast3; ++ const uint i02 = i12 / p.broadcast2; ++ ++ const uint batch_idx_a = i03 * p.ne02 + i02; ++#endif ++ ++ const uint blocks_m = (p.M + BM - 1) / BM; ++ const uint ir = gl_WorkGroupID.x % blocks_m; ++ const uint ik = gl_WorkGroupID.x / blocks_m; ++ const uint ic = gl_WorkGroupID.y; ++ ++ const uint WNITER = (WM * WN) / (WARP * TM * TN * WMITER); ++ const uint WSUBM = WM / WMITER; ++ const uint WSUBN = WN / WNITER; ++ ++#ifdef COOPMAT ++ const uint warp_i = gl_SubgroupID; ++ ++ const uint tiw = gl_SubgroupInvocationID; ++ ++ const uint cms_per_row = WM / TM; ++ const uint cms_per_col = WN / TN; ++ ++ const uint storestride = WARP / TM; ++ const uint store_r = tiw % TM; ++ const uint store_c = tiw / TM; ++#else ++ const uint warp_i = gl_LocalInvocationID.x / WARP; ++ ++ const uint tiw = gl_LocalInvocationID.x % WARP; ++ ++ const uint tiwr = tiw % (WSUBM / TM); ++ const uint tiwc = tiw / (WSUBM / TM); ++#endif ++ ++ const uint warp_r = warp_i % (BM / WM); ++ const uint warp_c = warp_i / (BM / WM); ++ ++ const uint loadr_a = gl_LocalInvocationID.x % (BK / LOAD_VEC_A); ++ const uint loadc_a = gl_LocalInvocationID.x / (BK / LOAD_VEC_A); ++ const uint loadr_b = gl_LocalInvocationID.x % (BK / LOAD_VEC_B); ++ const uint loadc_b = gl_LocalInvocationID.x / (BK / LOAD_VEC_B); ++ ++ const uint loadstride_a = gl_WorkGroupSize.x * LOAD_VEC_A / BK; ++ const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B / BK; ++ ++#ifdef MUL_MAT_ID ++ uint _ne1 = 0; ++ for (uint ii1 = 0; ii1 < p.nei1; ii1++) { ++ for (uint ii0 = 0; ii0 < p.nei0; ii0++) { ++ if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { ++ row_ids[_ne1] = u16vec2(ii0, ii1); ++ _ne1++; ++ } ++ } ++ } ++ ++ barrier(); ++ ++ // Workgroup has no work ++ if (ic * BN >= _ne1) return; ++#endif ++ ++#ifdef MUL_MAT_ID ++ const uint start_k = 0; ++ const uint end_k = p.K; ++#else ++ const uint start_k = ik * p.k_split; ++ const uint end_k = min(p.K, (ik + 1) * p.k_split); ++#endif ++ ++ uint pos_a = ( ++#ifdef MUL_MAT_ID ++ expert_idx * p.batch_stride_a + ++#else ++ batch_idx_a * p.batch_stride_a + ++#endif ++ ir * BM * p.stride_a + start_k) / LOAD_VEC_A; ++#ifdef MUL_MAT_ID ++ uint pos_b = 0; ++#else ++ uint pos_b = (batch_idx * p.batch_stride_b + ic * BN * p.stride_b + start_k) / LOAD_VEC_B; ++#endif ++ ++#ifdef COOPMAT ++ coopmat cache_a; ++ coopmat cache_b; ++ coopmat sums[cms_per_row * cms_per_col]; ++ ++ [[unroll]] for (uint i = 0; i < cms_per_row * cms_per_col; i++) { ++ sums[i] = coopmat(0.0f); ++ } ++#else ++ ACC_TYPE sums[WMITER * TM * WNITER * TN]; ++ FLOAT_TYPE cache_a[WMITER * TM]; ++ FLOAT_TYPE cache_b[WNITER * TN]; ++ ++ [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN; i++) { ++ sums[i] = ACC_TYPE(0.0f); ++ } ++#endif ++ ++ for (uint block = start_k; block < end_k; block += BK) { ++ [[unroll]] for (uint l = 0; l < BM; l += loadstride_a) { ++ ++#if defined(DATA_A_F32) || defined(DATA_A_F16) ++#if LOAD_VEC_A == 8 ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; ++ buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx][0].x); ++ buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx][0].y); ++ buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx][0].z); ++ buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx][0].w); ++ buf_a[buf_idx + 4] = FLOAT_TYPE(data_a[idx][1].x); ++ buf_a[buf_idx + 5] = FLOAT_TYPE(data_a[idx][1].y); ++ buf_a[buf_idx + 6] = FLOAT_TYPE(data_a[idx][1].z); ++ buf_a[buf_idx + 7] = FLOAT_TYPE(data_a[idx][1].w); ++#elif LOAD_VEC_A == 4 ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; ++ buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx].x); ++ buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx].y); ++ buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx].z); ++ buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx].w); ++#else ++ if (ir * BM + loadc_a + l < p.M && block + loadr_a < end_k) { ++ buf_a[(loadc_a + l) * SHMEM_STRIDE + loadr_a] = FLOAT_TYPE(data_a[pos_a + (loadc_a + l) * p.stride_a + loadr_a]); ++ } else { ++ buf_a[(loadc_a + l) * SHMEM_STRIDE + loadr_a] = FLOAT_TYPE(0.0f); ++ } ++#endif ++#elif defined(DATA_A_Q4_0) ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; ++ ++ const uint ib = idx / 16; ++ const uint iqs = idx & 0xF; ++ ++ const float d = float(data_a[ib].d); ++ const uint vui = uint(data_a[ib].qs[iqs]); ++ const vec2 v = (vec2(vui & 0xF, vui >> 4) - 8.0f) * d; ++ ++ buf_a[buf_idx ] = FLOAT_TYPE(v.x); ++ buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); ++#elif defined(DATA_A_Q4_1) ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; ++ ++ const uint ib = idx / 16; ++ const uint iqs = idx & 0xF; ++ ++ const float d = float(data_a[ib].d); ++ const float m = float(data_a[ib].m); ++ const uint vui = uint(data_a[ib].qs[iqs]); ++ const vec2 v = vec2(vui & 0xF, vui >> 4) * d + m; ++ ++ buf_a[buf_idx ] = FLOAT_TYPE(v.x); ++ buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); ++#elif defined(DATA_A_Q5_0) ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; ++ ++ const uint ib = idx / 16; ++ const uint iqs = idx & 0xF; ++ ++ const float d = float(data_a[ib].d); ++ const uint uint_qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0]; ++ const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); ++ const uint vui = uint(data_a[ib].qs[iqs]); ++ const vec2 v = (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f) * d; ++ ++ buf_a[buf_idx ] = FLOAT_TYPE(v.x); ++ buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); ++#elif defined(DATA_A_Q5_1) ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; ++ ++ const uint ib = idx / 16; ++ const uint iqs = idx & 0xF; ++ ++ const float d = float(data_a[ib].d); ++ const float m = float(data_a[ib].m); ++ const uint uint_qh = data_a[ib].qh; ++ const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); ++ const uint vui = uint(data_a[ib].qs[iqs]); ++ const vec2 v = vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) * d + m; ++ ++ buf_a[buf_idx ] = FLOAT_TYPE(v.x); ++ buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); ++#elif defined(DATA_A_Q8_0) ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; ++ ++ const uint ib = idx / 16; ++ const uint iqs = (idx & 0xF) * 2; ++ ++ const float d = float(data_a[ib].d); ++ const vec2 v = vec2(int(data_a[ib].qs[iqs]), int(data_a[ib].qs[iqs + 1])) * d; ++ ++ buf_a[buf_idx ] = FLOAT_TYPE(v.x); ++ buf_a[buf_idx + 1] = FLOAT_TYPE(v.y); ++#elif defined(DATA_A_Q2_K) ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; ++ ++ const uint ib = idx / 128; // 2 values per idx ++ const uint iqs = idx % 128; // 0..127 ++ ++ const uint qsi = (iqs / 64) * 32 + (iqs % 16) * 2; // 0,2,4..30 ++ const uint scalesi = iqs / 8; // 0..15 ++ const uint qsshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 ++ ++ const uvec2 qs = uvec2(data_a[ib].qs[qsi], data_a[ib].qs[qsi + 1]); ++ const uint scales = data_a[ib].scales[scalesi]; ++ const vec2 d = vec2(data_a[ib].d); ++ ++ const vec2 v = d.x * float(scales & 0xF) * vec2((qs >> qsshift) & 3) - d.y * float(scales >> 4); ++ ++ buf_a[buf_idx ] = FLOAT_TYPE(v.x); ++ buf_a[buf_idx + 1] = FLOAT_TYPE(v.y); ++#elif defined(DATA_A_Q3_K) ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; ++ ++ const uint ib = idx / 128; // 2 values per idx ++ const uint iqs = idx % 128; // 0..127 ++ ++ const uint n = iqs / 64; // 0,1 ++ const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 ++ const uint hmi = (iqs % 16) * 2; // 0,2,4..30 ++ const uint j = (iqs % 64) / 4; // 0..3 ++ const uint is = iqs / 8; // 0..15 ++ const uint halfsplit = ((iqs % 64) / 16); // 0,1,2,3 ++ const uint qsshift = halfsplit * 2; // 0,2,4,6 ++ const uint m = 1 << (4 * n + halfsplit); // 1,2,4,8,16,32,64,128 ++ ++ const int8_t us = int8_t(is < 4 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+8] >> 0) & 3) << 4) : ++ is < 8 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+4] >> 2) & 3) << 4) : ++ is < 12 ? (data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is+0] >> 4) & 3) << 4) : ++ (data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is-4] >> 6) & 3) << 4)); ++ const float dl = float(data_a[ib].d) * float(us - 32); ++ ++ buf_a[buf_idx ] = FLOAT_TYPE(dl * float(int8_t((data_a[ib].qs[qsi ] >> qsshift) & 3) - (((data_a[ib].hmask[hmi ] & m) != 0) ? 0 : 4))); ++ buf_a[buf_idx + 1] = FLOAT_TYPE(dl * float(int8_t((data_a[ib].qs[qsi + 1] >> qsshift) & 3) - (((data_a[ib].hmask[hmi + 1] & m) != 0) ? 0 : 4))); ++#elif defined(DATA_A_Q4_K) ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; ++ ++ const uint ib = idx / 128; // 2 values per idx ++ const uint iqs = idx % 128; // 0..127 ++ ++ const uint n = iqs / 32; // 0,1,2,3 ++ const uint b = (iqs % 32) / 16; // 0,1 ++ const uint is = 2 * n + b; // 0..7 ++ const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 ++ ++ const vec2 loadd = vec2(data_a[ib].d); ++ ++ const uint scidx0 = (is < 4) ? is : (is + 4); ++ const uint scidx1 = (is < 4) ? is : (is - 4); ++ const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ const uint scidxshift1 = (is < 4) ? 0 : 2; ++ const uint mbidx0 = is + 4; ++ const uint mbidx1 = (is < 4) ? is + 4 : is; ++ const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; ++ const uint mbidxshift0 = (is < 4) ? 0 : 4; ++ const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ const uint mbidxshift1 = (is < 4) ? 0 : 2; ++ ++ const uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); ++ const uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); ++ ++ const float d = loadd.x * sc; ++ const float m = -loadd.y * mbyte; ++ ++ buf_a[buf_idx ] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi ] >> (b * 4)) & 0xF), m)); ++ buf_a[buf_idx + 1] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi + 1] >> (b * 4)) & 0xF), m)); ++#elif defined(DATA_A_Q5_K) ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; ++ ++ const uint ib = idx / 128; // 2 values per idx ++ const uint iqs = idx % 128; // 0..127 ++ ++ const uint n = iqs / 32; // 0,1,2,3 ++ const uint b = (iqs % 32) / 16; // 0,1 ++ const uint is = 2 * n + b; // 0..7 ++ const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 ++ const uint qhi = (iqs % 16) * 2; // 0,2,4..30 ++ ++ const uint8_t hm = uint8_t(1 << (iqs / 16)); ++ ++ const vec2 loadd = vec2(data_a[ib].d); ++ ++ const uint scidx0 = (is < 4) ? is : (is + 4); ++ const uint scidx1 = (is < 4) ? is : (is - 4); ++ const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ const uint scidxshift1 = (is < 4) ? 0 : 2; ++ const uint mbidx0 = is + 4; ++ const uint mbidx1 = (is < 4) ? is + 4 : is; ++ const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; ++ const uint mbidxshift0 = (is < 4) ? 0 : 4; ++ const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; ++ const uint mbidxshift1 = (is < 4) ? 0 : 2; ++ ++ const uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); ++ const uint8_t mbyte = uint8_t(((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); ++ ++ const float d = loadd.x * sc; ++ const float m = -loadd.y * mbyte; ++ ++ buf_a[buf_idx ] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi ] >> (b * 4)) & 0xF) + float((data_a[ib].qh[qhi ] & hm) != 0 ? 16 : 0), m)); ++ buf_a[buf_idx + 1] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi + 1] >> (b * 4)) & 0xF) + float((data_a[ib].qh[qhi + 1] & hm) != 0 ? 16 : 0), m)); ++#elif defined(DATA_A_Q6_K) ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; ++ ++ const uint ib = idx / 128; // 2 values per idx ++ const uint iqs = idx % 128; // 0..127 ++ ++ const uint n = iqs / 64; // 0,1 ++ const uint b = (iqs % 64) / 32; // 0,1 ++ const uint is_b = (iqs % 16) / 8; // 0,1 ++ const uint qhshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 ++ const uint is = 8 * n + qhshift + is_b; // 0..15 ++ const uint qsi = n * 64 + (iqs % 32) * 2; // 0,2,4..126 ++ const uint qhi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 ++ ++ const float dscale = float(data_a[ib].d) * float(data_a[ib].scales[is]); ++ ++ buf_a[buf_idx ] = FLOAT_TYPE(dscale * float(int8_t(((data_a[ib].ql[qsi ] >> (b * 4)) & 0xF) | (((data_a[ib].qh[qhi ] >> qhshift) & 3) << 4)) - 32)); ++ buf_a[buf_idx + 1] = FLOAT_TYPE(dscale * float(int8_t(((data_a[ib].ql[qsi + 1] >> (b * 4)) & 0xF) | (((data_a[ib].qh[qhi + 1] >> qhshift) & 3) << 4)) - 32)); ++#elif defined(DATA_A_IQ4_NL) ++ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; ++ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; ++ ++ const uint ib = idx / 16; ++ const uint iqs = idx & 0xF; ++ ++ const float d = float(data_a[ib].d); ++ const uint vui = uint(data_a[ib].qs[iqs]); ++ const vec2 v = vec2(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[vui >> 4]) * d; ++ ++ buf_a[buf_idx ] = FLOAT_TYPE(v.x); ++ buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); ++#endif ++ } ++ [[unroll]] for (uint l = 0; l < BN; l += loadstride_b) { ++#if LOAD_VEC_B == 8 ++#ifdef MUL_MAT_ID ++ const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l]; ++ const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + loadr_b; ++#else ++ const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b; ++#endif ++ const uint buf_idx = (loadc_b + l) * SHMEM_STRIDE + loadr_b * LOAD_VEC_B; ++ buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx][0].x); ++ buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx][0].y); ++ buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx][0].z); ++ buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx][0].w); ++ buf_b[buf_idx + 4] = FLOAT_TYPE(data_b[idx][1].x); ++ buf_b[buf_idx + 5] = FLOAT_TYPE(data_b[idx][1].y); ++ buf_b[buf_idx + 6] = FLOAT_TYPE(data_b[idx][1].z); ++ buf_b[buf_idx + 7] = FLOAT_TYPE(data_b[idx][1].w); ++#elif LOAD_VEC_B == 4 ++#ifdef MUL_MAT_ID ++ const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l]; ++ const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + loadr_b; ++#else ++ const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b; ++#endif ++ const uint buf_idx = (loadc_b + l) * SHMEM_STRIDE + loadr_b * LOAD_VEC_B; ++ buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx].x); ++ buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx].y); ++ buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx].z); ++ buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx].w); ++#elif !MUL_MAT_ID ++ if (ic * BN + loadc_b + l < p.N && block + loadr_b < end_k) { ++ buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(data_b[pos_b + (loadc_b + l) * p.stride_b + loadr_b]); ++ } else { ++ buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(0.0f); ++ } ++#else ++ const uint row_i = ic * BN + loadc_b + l; ++ if (row_i < _ne1) { ++ const u16vec2 row_idx = row_ids[row_i]; ++ buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(data_b[pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + loadr_b]); ++ } else { ++ buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(0.0f); ++ } ++#endif ++ } ++ ++ barrier(); ++ ++ pos_a += BK / LOAD_VEC_A; ++ pos_b += BK / LOAD_VEC_B; ++ ++#ifdef COOPMAT ++ [[unroll]] for (uint i = 0; i < BK; i += TK) { ++ [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { ++ // Load from shared into cache ++ coopMatLoad(cache_a, buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor); ++ ++ [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { ++ coopMatLoad(cache_b, buf_b, (warp_c * WN + cm_col * TN) * SHMEM_STRIDE + i, SHMEM_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); ++ ++ sums[cm_col * cms_per_row + cm_row] = coopMatMulAdd(cache_a, cache_b, sums[cm_col * cms_per_row + cm_row]); ++ } ++ } ++ } ++#else ++ [[unroll]] for (uint i = 0; i < BK; i++) { ++ // Load from shared into cache ++ [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { ++ [[unroll]] for (uint j = 0; j < TM; j++) { ++ cache_a[wsir * TM + j] = buf_a[(warp_r * WM + wsir * WSUBM + tiwr * TM + j) * SHMEM_STRIDE + i]; ++ } ++ } ++ [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { ++ [[unroll]] for (uint j = 0; j < TN; j++) { ++ cache_b[wsic * TN + j] = buf_b[(warp_c * WN + wsic * WSUBN + tiwc * TN + j) * SHMEM_STRIDE + i]; ++ } ++ } ++ ++ [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { ++ [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { ++ [[unroll]] for (uint cc = 0; cc < TN; cc++) { ++ [[unroll]] for (uint cr = 0; cr < TM; cr++) { ++ const uint sums_idx = (wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr; ++ sums[sums_idx] = fma(ACC_TYPE(cache_a[wsir * TM + cr]), ACC_TYPE(cache_b[wsic * TN + cc]), sums[sums_idx]); ++ } ++ } ++ } ++ } ++ } ++#endif ++ ++ barrier(); ++ } ++ ++ const uint dr = ir * BM + warp_r * WM; ++ const uint dc = ic * BN + warp_c * WN; ++ ++#ifndef MUL_MAT_ID ++ const uint offsets = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * gl_NumWorkGroups.z; ++#endif ++ ++#ifdef COOPMAT ++#ifdef MUL_MAT_ID ++ [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { ++ [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { ++ coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); ++ ++ [[unroll]] for (uint col = 0; col < BN; col += storestride) { ++ const uint row_i = dc + cm_col * TN + col + store_c; ++ if (row_i >= _ne1) break; ++ ++ const u16vec2 row_idx = row_ids[row_i]; ++ ++ data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); ++ } ++ } ++ } ++#else ++ const bool is_aligned = p.stride_d % 4 == 0; // Assumption: D_TYPE == float ++ ++ [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { ++ [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { ++ const bool is_in_bounds = dr + (cm_row + 1) * TM <= p.M && dc + (cm_col + 1) * TN <= p.N; ++ ++ if (is_aligned && is_in_bounds) { ++ // Full coopMat is within bounds and stride_d is aligned with 16B ++ coopmat cm_dtype = coopmat(sums[cm_col * cms_per_row + cm_row]); ++ coopMatStore(cm_dtype, data_d, offsets + (dc + cm_col * TN) * p.stride_d + dr + cm_row * TM, p.stride_d, gl_CooperativeMatrixLayoutColumnMajor); ++ } else if (is_in_bounds) { ++ // Full coopMat is within bounds, but stride_d is not aligned ++ coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); ++ ++ [[unroll]] for (uint col = 0; col < TN; col += storestride) { ++ data_d[offsets + (dc + cm_col * TN + col + store_c) * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); ++ } ++ } else if (dr + cm_row * TM < p.M && dc + cm_col * TN < p.N) { ++ // Partial coopMat is within bounds ++ coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); ++ ++ [[unroll]] for (uint col = 0; col < TN; col += storestride) { ++ if (dr + cm_row * TM + store_r < p.M && dc + cm_col * TN + col + store_c < p.N) { ++ data_d[offsets + (dc + cm_col * TN + col + store_c) * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); ++ } ++ } ++ } ++ } ++ } ++#endif // MUL_MAT_ID ++#else ++ [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { ++ [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { ++ ++ const uint dr_warp = dr + wsir * WSUBM + tiwr * TM; ++ const uint dc_warp = dc + wsic * WSUBN + tiwc * TN; ++ [[unroll]] for (uint cc = 0; cc < TN; cc++) { ++#ifdef MUL_MAT_ID ++ const uint row_i = dc_warp + cc; ++ if (row_i >= _ne1) break; ++ ++ const u16vec2 row_idx = row_ids[row_i]; ++#endif // MUL_MAT_ID ++ [[unroll]] for (uint cr = 0; cr < TM; cr++) { ++#ifdef MUL_MAT_ID ++ data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]); ++#else ++ if (dr_warp + cr < p.M && dc_warp + cc < p.N) { ++ data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]); ++ } ++#endif // MUL_MAT_ID ++ } ++ } ++ } ++ } ++#endif // COOPMAT ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp +new file mode 100644 +index 00000000..cbfa5dce +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp +@@ -0,0 +1,328 @@ ++#version 450 ++ ++#extension GL_EXT_control_flow_attributes : enable ++#extension GL_EXT_shader_16bit_storage : require ++ ++#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require ++#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require ++#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require ++#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require ++ ++#extension GL_KHR_memory_scope_semantics : enable ++#extension GL_KHR_cooperative_matrix : enable ++#extension GL_NV_cooperative_matrix2 : enable ++#extension GL_EXT_buffer_reference : enable ++#extension GL_KHR_shader_subgroup_ballot : enable ++#extension GL_KHR_shader_subgroup_vote : enable ++ ++#include "types.comp" ++ ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++layout (constant_id = 1) const uint BM = 64; ++layout (constant_id = 2) const uint BN = 64; ++layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working with a quant ++ ++layout (push_constant) uniform parameter ++{ ++ uint M; ++ uint N; ++ uint K; ++ uint stride_a; ++ uint stride_b; ++ uint stride_d; ++ ++ uint batch_stride_a; ++ uint batch_stride_b; ++ uint batch_stride_d; ++ ++#ifdef MUL_MAT_ID ++ uint nei0; ++ uint nei1; ++ uint nbi1; ++ uint ne11; ++#else ++ uint k_split; ++ uint ne02; ++ uint ne12; ++ uint broadcast2; ++ uint broadcast3; ++#endif ++} p; ++ ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; ++layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; ++ ++#if QUANT_K > 1 ++#define DECODEFUNCA , dequantFuncA ++#define MAT_A_TYPE float16_t ++ ++#include "dequant_funcs_cm2.comp" ++ ++#else ++#define DECODEFUNCA ++#define MAT_A_TYPE A_TYPE ++#endif ++ ++#define MAT_B_TYPE B_TYPE ++ ++#ifdef MUL_MAT_ID ++layout (binding = 3) readonly buffer IDS {int data_ids[];}; ++ ++shared u16vec4 row_ids[3072]; ++ ++layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufB { ++ B_TYPE b[]; ++}; ++ ++uint _ne1; ++shared uint _ne1_sh; ++ ++B_TYPE decodeFuncB(const in decodeBufB bl, const in uint blockCoords[2], const in uint coordInBlock[2]) ++{ ++ const uint row_i = blockCoords[0]; ++ ++ if (row_i >= _ne1) { ++ return B_TYPE(0.0); ++ } ++ ++ const u16vec4 row_idx = row_ids[row_i]; ++ B_TYPE ret = data_b[row_idx.y * p.batch_stride_b + row_idx.x * p.stride_b + blockCoords[1]]; ++ ++ return ret; ++} ++ ++D_TYPE perElemOpD(const in uint32_t r, const in uint32_t c, const in D_TYPE elem, const in uint32_t ir, const in uint32_t ic) ++{ ++ uint dr = ir * BM + r; ++ uint dc = ic * BN + c; ++ ++ if (dr < p.M && dc < _ne1) { ++ uint row_i = dc; ++ const u16vec4 row_idx = row_ids[row_i]; ++ data_d[row_idx.y * p.batch_stride_d + row_idx.z * p.stride_d + dr] = elem; ++ } ++ return elem; ++} ++ ++#endif ++ ++void main() { ++#if defined(DATA_A_IQ4_NL) ++ init_iq4nl_shmem(); ++#endif ++ ++#ifdef MUL_MAT_ID ++ const uint expert_idx = gl_GlobalInvocationID.z; ++#else ++ const uint batch_idx = gl_GlobalInvocationID.z; ++ ++ const uint i13 = batch_idx / p.ne12; ++ const uint i12 = batch_idx % p.ne12; ++ ++ const uint i03 = i13 / p.broadcast3; ++ const uint i02 = i12 / p.broadcast2; ++ ++ const uint batch_idx_a = i03 * p.ne02 + i02; ++#endif ++ ++ const uint blocks_m = (p.M + BM - 1) / BM; ++ const uint ir = gl_WorkGroupID.x % blocks_m; ++ const uint ik = gl_WorkGroupID.x / blocks_m; ++ const uint ic = gl_WorkGroupID.y; ++ ++#ifdef MUL_MAT_ID ++ // Spread the search across all elements in the first subgroup ++ if (gl_SubgroupID == 0) { ++ _ne1 = 0; ++ uint num_elements = p.nei1 * p.nei0; ++ ++ for (uint i = gl_SubgroupInvocationID; subgroupAny(i < num_elements); i += gl_SubgroupSize) { ++ bool in_range = i < num_elements; ++ uint ii0 = i % p.nei0; ++ uint ii1 = i / p.nei0; ++ uint id = in_range ? data_ids[ii1*p.nbi1 + ii0] : 0; ++ uvec4 ballot = subgroupBallot(in_range && id == expert_idx); ++ uint idx = subgroupBallotExclusiveBitCount(ballot); ++ if (in_range && id == expert_idx) { ++ row_ids[_ne1 + idx] = u16vec4(ii0 % p.ne11, ii1, ii0, 0); ++ } ++ _ne1 += subgroupBallotBitCount(ballot); ++ } ++ _ne1_sh = _ne1; ++ } ++ ++ barrier(); ++ ++ _ne1 = _ne1_sh; ++ ++ // Workgroup has no work ++ if (ic * BN >= _ne1) return; ++#endif ++ ++#ifdef MUL_MAT_ID ++ uint start_k = 0; ++ const uint end_k = p.K; ++#else ++ uint start_k = ik * p.k_split; ++ const uint end_k = min(p.K, (ik + 1) * p.k_split); ++#endif ++ ++ coopmat sum; ++ sum = coopmat(0.0); ++ ++#ifdef MUL_MAT_ID ++ uint pos_a = (expert_idx * p.batch_stride_a) / QUANT_K; ++ uint pos_b = 0; ++#else ++ uint pos_a = (batch_idx_a * p.batch_stride_a) / QUANT_K; ++ uint pos_b = batch_idx * p.batch_stride_b; ++#endif ++ ++ uint stride_a = p.stride_a / QUANT_K; ++ uint stride_b = p.stride_b; ++ ++ // Hint to the compiler that values are aligned (want 16B alignment). ++ // Quants are always block-aligned, no alignment needed. ++#if ALIGNED ++#if QUANT_K == 1 ++ stride_a &= ~7; ++#endif ++ stride_b &= ~7; ++#endif ++ ++ // Create layouts for both clamped and unclamped accesses ++ tensorLayoutNV<2> tensorLayoutA = createTensorLayoutNV(2); ++ tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutAClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); ++ tensorLayoutNV<2> tensorLayoutB = createTensorLayoutNV(2); ++ tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutBClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); ++ tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); ++ ++#if QUANT_K > 1 ++ tensorLayoutA = setTensorLayoutBlockSizeNV(tensorLayoutA, 1, QUANT_K); ++ tensorLayoutAClamp = setTensorLayoutBlockSizeNV(tensorLayoutAClamp, 1, QUANT_K); ++#endif ++ ++ // Use end_k rather than p.K as the dimension because that's what ++ // we need to bound check against when using split_k ++ tensorLayoutA = setTensorLayoutDimensionNV(tensorLayoutA, p.M, end_k); ++ tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, p.N, end_k); ++ tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.N, p.M); ++ tensorLayoutAClamp = setTensorLayoutDimensionNV(tensorLayoutAClamp, p.M, end_k); ++ tensorLayoutBClamp = setTensorLayoutDimensionNV(tensorLayoutBClamp, p.N, end_k); ++ ++ tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0); ++ ++#if !defined(MUL_MAT_ID) ++ // Detect a fast path where all loads are entirely in bounds and no clamping is required ++ if ((ir + 1) * BM <= p.M && (ic + 1) * BN <= p.N && (start_k % BK) == 0 && (end_k % BK) == 0 && ++#if QUANT_K == 1 ++ (stride_a % 8) == 0 && ++#endif ++ (stride_b % 8) == 0 && (start_k % 8) == 0) { ++ // Hint to the compiler that values are aligned (want 16B alignment) ++ start_k &= ~7; ++ stride_b &= ~7; ++#if QUANT_K == 1 ++ stride_a &= ~7; ++#endif ++ ++ tensorLayoutA = setTensorLayoutStrideNV(tensorLayoutA, stride_a, 1); ++ tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1); ++ ++ uint k_iters = (end_k - start_k + BK - 1) / BK; ++ ++ for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) { ++ ++ coopmat mat_a; ++ coopmat mat_b; ++ ++ coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA); ++ coopmat mat_a_ft = coopmat(mat_a); ++ ++ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose); ++ coopmat mat_b_ft = coopmat(mat_b); ++ ++ sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); ++ } ++ } else ++#endif // !defined(MUL_MAT_ID) ++ { ++ tensorLayoutA = setTensorLayoutStrideNV(tensorLayoutA, stride_a, 1); ++ ++ tensorLayoutAClamp = setTensorLayoutStrideNV(tensorLayoutAClamp, stride_a, 1); ++ ++ tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1); ++ ++ tensorLayoutBClamp = setTensorLayoutStrideNV(tensorLayoutBClamp, stride_b, 1); ++ ++ [[dont_unroll]] ++ for (uint block_k = start_k; block_k < end_k; block_k += BK) { ++ ++ coopmat mat_a; ++ coopmat mat_b; ++ coopmat mat_a_ft; ++ coopmat mat_b_ft; ++ ++ // Clamping is expensive, so detect different code paths for each combination ++ // of A and B needing clamping. ++ bool unclampedA = (ir + 1) * BM <= p.M && block_k + BK <= end_k && (block_k % 8) == 0; ++#ifdef MUL_MAT_ID ++ bool unclampedB = true; ++#else ++ bool unclampedB = (ic + 1) * BN <= p.N && block_k + BK <= end_k && (block_k % 8) == 0; ++#endif ++ if (unclampedA && unclampedB) { ++ coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, (block_k & ~7), BK) DECODEFUNCA); ++#ifdef MUL_MAT_ID ++ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose, decodeFuncB); ++#else ++ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, (block_k & ~7), BK), tensorViewTranspose); ++#endif ++ mat_a_ft = coopmat(mat_a); ++ mat_b_ft = coopmat(mat_b); ++ sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); ++ } else if (unclampedA && !unclampedB) { ++ coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, (block_k & ~7), BK) DECODEFUNCA); ++ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutBClamp, ic * BN, BN, block_k, BK), tensorViewTranspose); ++ ++ mat_a_ft = coopmat(mat_a); ++ mat_b_ft = coopmat(mat_b); ++ sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); ++ } else if (!unclampedA && unclampedB) { ++ coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutAClamp, ir * BM, BM, block_k, BK) DECODEFUNCA); ++#ifdef MUL_MAT_ID ++ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose, decodeFuncB); ++#else ++ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, (block_k & ~7), BK), tensorViewTranspose); ++#endif ++ mat_a_ft = coopmat(mat_a); ++ mat_b_ft = coopmat(mat_b); ++ sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); ++ } else if (!unclampedA && !unclampedB) { ++ coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutAClamp, ir * BM, BM, block_k, BK) DECODEFUNCA); ++ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutBClamp, ic * BN, BN, block_k, BK), tensorViewTranspose); ++ ++ mat_a_ft = coopmat(mat_a); ++ mat_b_ft = coopmat(mat_b); ++ sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); ++ } ++ } ++ } ++ ++ // Convert from ACC_TYPE to D_TYPE ++ coopmat mat_d; ++ mat_d = coopmat(sum); ++ ++#ifdef MUL_MAT_ID ++ // Call callback to store each element, remapping row through shared memory ++ coopMatPerElementNV(mat_d, mat_d, perElemOpD, ir, ic); ++#else ++ tensorLayoutD = setTensorLayoutStrideNV(tensorLayoutD, p.stride_d, 1); ++ ++ uint pos_d = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * gl_NumWorkGroups.z; ++ coopMatStoreTensorNV(mat_d, data_d, pos_d, sliceTensorLayoutNV(tensorLayoutD, ic * BN, BN, ir * BM, BM), tensorViewTranspose); ++#endif ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp +new file mode 100644 +index 00000000..6627a50b +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp +@@ -0,0 +1,44 @@ ++#version 450 ++ ++#include "generic_head.comp" ++#include "types.comp" ++ ++#extension GL_EXT_control_flow_attributes : enable ++#define BLOCK_SIZE 512 ++ ++layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++shared vec2 sum[BLOCK_SIZE]; ++ ++void main() { ++ const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; ++ const uint tid = gl_LocalInvocationID.x; ++ ++ sum[tid] = vec2(0.0f, 0.0f); ++ ++ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { ++ const float xi = float(data_a[row*p.KX + col]); ++ sum[tid].x += xi; ++ sum[tid].y += xi * xi; ++ } ++ ++ // sum up partial sums and write back result ++ barrier(); ++ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { ++ if (tid < s) { ++ sum[tid] += sum[tid + s]; ++ } ++ barrier(); ++ } ++ ++ const float mean = sum[0].x / p.KX; ++ const float var = sum[0].y / p.KX - mean * mean; ++ const float inv_std = inversesqrt(var + p.param1); ++ ++ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { ++ data_d[row*p.KX + col] = D_TYPE((float(data_a[row*p.KX + col]) - mean) * inv_std); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp +new file mode 100644 +index 00000000..450b67fc +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp +@@ -0,0 +1,28 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_unary_head.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; ++ ++ if (idx >= p.ne) { ++ return; ++ } ++ ++ const uint i3 = idx / (p.ne12*p.ne11*p.ne10); ++ const uint i3_offset = i3 * p.ne12*p.ne11*p.ne10; ++ const uint i2 = (idx - i3_offset) / (p.ne11*p.ne10); ++ const uint i2_offset = i2*p.ne11*p.ne10; ++ const uint i1 = (idx - i3_offset - i2_offset) / p.ne10; ++ const uint i0 = idx - i3_offset - i2_offset - i1*p.ne10; ++ ++ const uint src0_idx = i3*p.nb03 + i2*p.nb02 + i1*p.nb01 + i0*p.nb00; ++ const uint dst_idx = i3*p.nb13 + i2*p.nb12 + i1*p.nb11 + i0*p.nb10; ++ ++ const bool is_src0 = i0 < p.ne00 && i1 < p.ne01 && i2 < p.ne02 && i3 < p.ne03; ++ ++ data_d[get_doffset() + dst_idx] = D_TYPE(is_src0 ? data_a[get_aoffset() + src0_idx] : 0.0f); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp +new file mode 100644 +index 00000000..b6124411 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp +@@ -0,0 +1,74 @@ ++#version 450 ++ ++#include "types.comp" ++ ++#extension GL_EXT_shader_16bit_storage : require ++ ++layout(push_constant) uniform parameter { ++ uint IW; uint IH; ++ uint OW; uint OH; ++ uint OC; ++ uint pelements; ++ uint op; ++ int k0; int k1; ++ int s0; int s1; ++ int p0; int p1; ++} p; ++ ++#define BLOCK_SIZE 512 ++#define FLT_MAX 3.402823466e+38F ++#define OP_POOL_MAX 0u ++#define OP_POOL_AVG 1u ++ ++layout (local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; ++ ++layout(binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout(binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++void main() { ++ const uint idx = gl_GlobalInvocationID.x; ++ if (idx >= p.pelements) { ++ return; ++ } ++ ++ const uint O_HW = p.OW * p.OH; ++ ++ const uint nc = idx / O_HW; ++ const uint cur_oh = (idx % O_HW) / p.OW; ++ const uint cur_ow = (idx % O_HW) % p.OW; ++ ++ const int start_h = int(cur_oh) * p.s0 - p.p0; ++ const uint bh = max(start_h, 0); ++ const uint eh = min(start_h + p.k0, p.IH); ++ ++ const int start_w = int(cur_ow) * p.s1 - p.p1; ++ const uint bw = max(start_w, 0); ++ const uint ew = min(start_w + p.k1, p.IW); ++ ++ const float scale = 1.0 / float(p.k0 * p.k1); ++ float res; ++ ++ if (p.op == OP_POOL_AVG) { ++ res = 0.0; ++ } else if (p.op == OP_POOL_MAX) { ++ res = -FLT_MAX; ++ } else { ++ return; ++ } ++ ++ #pragma unroll ++ for (uint i = bh; i < eh; i++) { ++ #pragma unroll ++ for (uint j = bw; j < ew; j++) { ++ const float cur = D_TYPE(data_a[nc * p.IH * p.IW + i * p.IW + j]); ++ ++ if (p.op == OP_POOL_AVG) { ++ res += cur * scale; ++ } else if (p.op == OP_POOL_MAX) { ++ res = max(res, cur); ++ } ++ } ++ } ++ ++ data_d[nc * O_HW + cur_oh * p.OW + cur_ow] = res; ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp +new file mode 100644 +index 00000000..52a19b62 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp +@@ -0,0 +1,21 @@ ++#version 450 ++ ++#include "generic_head.comp" ++#include "types.comp" ++ ++#extension GL_EXT_control_flow_attributes : enable ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++void main() { ++ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; ++ ++ if (i >= p.KX) { ++ return; ++ } ++ ++ data_d[i] = max(float(data_a[i]), 0); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp +new file mode 100644 +index 00000000..1568b141 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp +@@ -0,0 +1,26 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_unary_head.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++uint src0_idx_mod(uint idx) { ++ const uint i13 = idx / (p.ne12*p.ne11*p.ne10); ++ const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10; ++ const uint i12 = (idx - i13_offset) / (p.ne11*p.ne10); ++ const uint i12_offset = i12*p.ne11*p.ne10; ++ const uint i11 = (idx - i13_offset - i12_offset) / p.ne10; ++ const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10; ++ return (i13 % p.ne03)*p.nb03 + (i12 % p.ne02)*p.nb02 + (i11 % p.ne01)*p.nb01 + (i10 % p.ne00)*p.nb00; ++} ++ ++void main() { ++ const uint idx = get_idx(); ++ ++ if (idx >= p.ne) { ++ return; ++ } ++ ++ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(data_a[get_aoffset() + src0_idx_mod(idx)]); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp +new file mode 100644 +index 00000000..b554400b +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp +@@ -0,0 +1,42 @@ ++#version 450 ++ ++#include "generic_head.comp" ++#include "types.comp" ++ ++#extension GL_EXT_control_flow_attributes : enable ++#define BLOCK_SIZE 512 ++ ++layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++shared FLOAT_TYPE sum[BLOCK_SIZE]; ++ ++void main() { ++ const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; ++ const uint tid = gl_LocalInvocationID.x; ++ ++ sum[tid] = FLOAT_TYPE(0.0f); // partial sum for thread in warp ++ ++ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { ++ const FLOAT_TYPE xi = FLOAT_TYPE(data_a[row*p.KX + col]); ++ sum[tid] += xi * xi; ++ } ++ ++ // sum up partial sums and write back result ++ barrier(); ++ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { ++ if (tid < s) { ++ sum[tid] += sum[tid + s]; ++ } ++ barrier(); ++ } ++ ++ const FLOAT_TYPE mean = sum[0] / FLOAT_TYPE(p.KX); ++ const FLOAT_TYPE scale = inversesqrt(mean + FLOAT_TYPE(p.param1)); ++ ++ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { ++ data_d[row*p.KX + col] = D_TYPE(scale * FLOAT_TYPE(data_a[row*p.KX + col])); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp +new file mode 100644 +index 00000000..574b51ca +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp +@@ -0,0 +1,49 @@ ++#include "types.comp" ++ ++#extension GL_EXT_shader_16bit_storage : require ++#extension GL_EXT_spirv_intrinsics: enable ++ ++#if RTE16 ++spirv_execution_mode(capabilities = [4467], 4462, 16); // RoundingModeRTE, 16 bits ++#endif ++ ++layout(local_size_x = 1, local_size_y = 256, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) readonly buffer Y {int data_pos[];}; ++layout (binding = 2) readonly buffer Z {float data_ff[];}; ++layout (binding = 3) writeonly buffer D {D_TYPE data_d[];}; ++ ++layout (push_constant) uniform parameter { ++ uint ncols; ++ uint n_dims; ++ float freq_scale; ++ uint p_delta_rows; ++ float freq_base; ++ float ext_factor; ++ float attn_factor; ++ float corr_dims[2]; ++ float theta_scale; ++ uint has_ff; ++} p; ++ ++float rope_yarn_ramp(const float low, const float high, const uint i0) { ++ const float y = (i0 / 2 - low) / max(0.001f, high - low); ++ return 1.0f - min(1.0f, max(0.0f, y)); ++} ++ ++void rope_yarn(const float theta_extrap, const uint i0, out float cos_theta, out float sin_theta) { ++ float mscale = p.attn_factor; ++ // Get n-d rotational scaling corrected for extrapolation ++ float theta_interp = p.freq_scale * theta_extrap; ++ float theta = theta_interp; ++ if (p.ext_factor != 0.0f) { ++ float ramp_mix = rope_yarn_ramp(p.corr_dims[0], p.corr_dims[1], i0) * p.ext_factor; ++ theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix; ++ ++ // Get n-d magnitude scaling corrected for interpolation ++ mscale *= 1.0f + 0.1f * log(1.0f / p.freq_scale); ++ } ++ cos_theta = cos(theta) * mscale; ++ sin_theta = sin(theta) * mscale; ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp +new file mode 100644 +index 00000000..83b46b69 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp +@@ -0,0 +1,37 @@ ++#version 450 ++ ++#include "rope_head.comp" ++ ++void main() { ++ const uint col = gl_GlobalInvocationID.y * 2; ++ const uint row = gl_GlobalInvocationID.x; ++ ++ if (col >= p.ncols) { ++ return; ++ } ++ ++ if (col >= p.n_dims) { ++ const uint i = row*p.ncols + col; ++ ++ data_d[i + 0] = data_a[i + 0]; ++ data_d[i + 1] = data_a[i + 1]; ++ ++ return; ++ } ++ ++ const uint i = row*p.ncols + col/2; ++ const uint i2 = row/p.p_delta_rows; ++ ++ const float theta_base = data_pos[i2] * pow(p.theta_scale, col/2.0f); ++ ++ const float freq_factor = p.has_ff != 0 ? data_ff[col/2] : 1.0f; ++ ++ float cos_theta, sin_theta; ++ rope_yarn(theta_base / freq_factor, col, cos_theta, sin_theta); ++ ++ const float x0 = float(data_a[i + 0]); ++ const float x1 = float(data_a[i + p.n_dims/2]); ++ ++ data_d[i + 0] = D_TYPE(x0*cos_theta - x1*sin_theta); ++ data_d[i + p.n_dims/2] = D_TYPE(x0*sin_theta + x1*cos_theta); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp +new file mode 100644 +index 00000000..e416ad93 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp +@@ -0,0 +1,37 @@ ++#version 450 ++ ++#include "rope_head.comp" ++ ++void main() { ++ const uint col = gl_GlobalInvocationID.y * 2; ++ const uint row = gl_GlobalInvocationID.x; ++ ++ if (col >= p.ncols) { ++ return; ++ } ++ ++ if (col >= p.n_dims) { ++ const uint i = row*p.ncols + col; ++ ++ data_d[i + 0] = data_a[i + 0]; ++ data_d[i + 1] = data_a[i + 1]; ++ ++ return; ++ } ++ ++ const uint i = row*p.ncols + col; ++ const uint i2 = row/p.p_delta_rows; ++ ++ const float theta_base = data_pos[i2] * pow(p.theta_scale, col/2.0f); ++ ++ const float freq_factor = p.has_ff != 0 ? data_ff[col/2] : 1.0f; ++ ++ float cos_theta, sin_theta; ++ rope_yarn(theta_base / freq_factor, col, cos_theta, sin_theta); ++ ++ const float x0 = float(data_a[i + 0]); ++ const float x1 = float(data_a[i + 1]); ++ ++ data_d[i + 0] = D_TYPE(x0*cos_theta - x1*sin_theta); ++ data_d[i + 1] = D_TYPE(x0*sin_theta + x1*cos_theta); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp +new file mode 100644 +index 00000000..4663428d +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp +@@ -0,0 +1,24 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_unary_head.comp" ++ ++const uint num_threads = 128; ++ ++layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ uint idx = get_idx(); ++ ++ // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation ++ const uint num_iter = 4; ++ ++ [[unroll]] for (uint i = 0; i < num_iter; ++i) { ++ if (idx >= p.ne) { ++ continue; ++ } ++ ++ data_d[get_doffset() + idx] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + idx]) * FLOAT_TYPE(p.param1)); ++ idx += num_threads; ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp +new file mode 100644 +index 00000000..4d36f88e +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp +@@ -0,0 +1,22 @@ ++#version 450 ++ ++#include "generic_head.comp" ++#include "types.comp" ++ ++#extension GL_EXT_control_flow_attributes : enable ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++void main() { ++ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; ++ ++ if (i >= p.KX) { ++ return; ++ } ++ ++ const float xi = float(data_a[i]); ++ data_d[i] = D_TYPE(xi / (1.0f + exp(-xi))); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp +new file mode 100644 +index 00000000..d7c15a16 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp +@@ -0,0 +1,17 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_unary_head.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ const uint idx = get_idx(); ++ ++ if (idx >= p.ne) { ++ return; ++ } ++ ++ const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); ++ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(sin(val)); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp +new file mode 100644 +index 00000000..a25808e1 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp +@@ -0,0 +1,174 @@ ++#version 450 ++ ++#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require ++#extension GL_EXT_control_flow_attributes : enable ++ ++layout (push_constant) uniform parameter ++{ ++ uint KX; ++ uint KY; ++ float scale; ++ float max_bias; ++ float m0; ++ float m1; ++ uint n_head_log2; ++ uint nrows_x; ++} p; ++ ++#include "types.comp" ++ ++layout(constant_id = 0) const uint BLOCK_SIZE = 32; ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) readonly buffer Y {B_TYPE data_b[];}; ++layout (binding = 2) buffer D {D_TYPE data_d[];}; ++ ++shared FLOAT_TYPE vals[BLOCK_SIZE]; ++ ++// num_iters is the number of BLOCK_SIZE loop iterations we need to iterate ++// over all the columns. The main function tries to pass a constant here, ++// as if it were a template function, to allow unrolling. ++void soft_max(uint num_iters) { ++ const uint tid = gl_LocalInvocationID.x; ++ const uint rowx = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; ++ const uint rowy = (p.KY > 0) ? (rowx % p.KY) : 0; ++ ++ if (rowx >= p.nrows_x) { ++ return; ++ } ++ ++ float slope = 1.0f; ++ ++ // ALiBi ++ if (p.max_bias > 0.0f) { ++ const uint h = rowx/p.KY; // head index ++ ++ const float base = h < p.n_head_log2 ? p.m0 : p.m1; ++ const uint exp = h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1; ++ ++ slope = pow(base, exp); ++ } ++ ++ // Find max ++ FLOAT_TYPE max_val = uintBitsToFloat(0xFF800000); ++ ++ // Cache values while we compute the max, so we don't need to read them ++ // again when we're ready to compute exp(x-max). ++ const uint DATA_CACHE_SIZE = 16; ++ FLOAT_TYPE data_cache[DATA_CACHE_SIZE]; ++ ++ [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { ++ const uint col = col0 + tid; ++ ++ FLOAT_TYPE a = FLOAT_TYPE(0); ++ if (col < p.KX) { ++ a = data_a[rowx * p.KX + col]; ++ } ++ ++ FLOAT_TYPE b = FLOAT_TYPE(0); ++ if (p.KY > 0 && col < p.KX) { ++ b = data_b[rowy * p.KX + col]; ++ } ++ ++ FLOAT_TYPE v = a * p.scale + slope * b; ++ ++ if (col < p.KX) { ++ max_val = max(max_val, v); ++ } ++ ++ if (idx < DATA_CACHE_SIZE) { ++ data_cache[idx] = v; ++ } ++ } ++ ++ // reduce across the workgroup ++ vals[tid] = max_val; ++ barrier(); ++ [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { ++ if (tid < s) { ++ vals[tid] = max(vals[tid], vals[tid + s]); ++ } ++ barrier(); ++ } ++ ++ max_val = vals[0]; ++ barrier(); ++ ++ FLOAT_TYPE sum = FLOAT_TYPE(0.0f); ++ ++ // Compute sum{exp(x - max)} ++ [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { ++ const uint col = col0 + tid; ++ ++ if (col >= p.KX) { ++ break; ++ } ++ ++ // compute exp(a*scale+b*slope), add it to sum, and cache the new value ++ // in data_cache if possible. ++ const uint i = rowx * p.KX + col; ++ FLOAT_TYPE val; ++ if (idx < DATA_CACHE_SIZE) { ++ val = exp(data_cache[idx] - max_val); ++ } else { ++ val = exp(FLOAT_TYPE(data_a[i]) * p.scale + (p.KY > 0 ? slope * FLOAT_TYPE(data_b[rowy * p.KX + col]) : FLOAT_TYPE(0.0f)) - max_val); ++ } ++ sum += val; ++ if (idx < DATA_CACHE_SIZE) { ++ data_cache[idx] = val; ++ } else { ++ data_d[i] = D_TYPE(val); ++ } ++ } ++ ++ // reduce across the workgroup ++ vals[tid] = sum; ++ barrier(); ++ [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { ++ if (tid < s) { ++ vals[tid] += vals[tid + s]; ++ } ++ barrier(); ++ } ++ sum = vals[0]; ++ ++ FLOAT_TYPE rcpdivisor = 1.0/sum; ++ ++ [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { ++ const uint col = col0 + tid; ++ ++ if (col >= p.KX) { ++ continue; ++ } ++ ++ if (idx < DATA_CACHE_SIZE) { ++ data_d[rowx*p.KX + col] = D_TYPE(data_cache[idx] * rcpdivisor); ++ } else { ++ data_d[rowx*p.KX + col] *= D_TYPE(rcpdivisor); ++ } ++ } ++} ++ ++void main() { ++ // instantiate the soft_max function for several different ++ // dimensions, to allow loop unrolling ++ uint num_blocks = (p.KX + BLOCK_SIZE - 1) / BLOCK_SIZE; ++ if (num_blocks > 32) { ++ soft_max(num_blocks); ++ } else if (num_blocks > 16) { ++ soft_max(32); ++ } else if (num_blocks > 8) { ++ soft_max(16); ++ } else if (num_blocks > 4) { ++ soft_max(8); ++ } else if (num_blocks == 4) { ++ soft_max(4); ++ } else if (num_blocks == 3) { ++ soft_max(3); ++ } else if (num_blocks == 2) { ++ soft_max(2); ++ } else if (num_blocks == 1) { ++ soft_max(1); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp +new file mode 100644 +index 00000000..ef43598b +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp +@@ -0,0 +1,17 @@ ++#version 450 ++ ++#include "types.comp" ++#include "generic_unary_head.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++void main() { ++ const uint idx = get_idx(); ++ ++ if (idx >= p.ne) { ++ return; ++ } ++ ++ const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); ++ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(val * val); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp +new file mode 100644 +index 00000000..961e5ffa +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp +@@ -0,0 +1,37 @@ ++#version 450 ++ ++#include "generic_head.comp" ++#include "types.comp" ++ ++#extension GL_EXT_control_flow_attributes : enable ++layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++layout (constant_id = 0) const uint BLOCK_SIZE = 32; ++ ++shared FLOAT_TYPE tmp[BLOCK_SIZE]; ++ ++void main() { ++ const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; ++ const uint col = gl_LocalInvocationID.x; ++ ++ tmp[col] = FLOAT_TYPE(0.0f); ++ ++ for (uint i = col; i < p.KX; i += BLOCK_SIZE) { ++ tmp[col] += FLOAT_TYPE(data_a[row*p.KX + i]); ++ } ++ ++ barrier(); ++ [[unroll]] for (int s = int(BLOCK_SIZE) / 2; s > 0; s >>= 1) { ++ if (col < s) { ++ tmp[col] += tmp[col + s]; ++ } ++ barrier(); ++ } ++ ++ if (col == 0) { ++ data_d[row] = D_TYPE(tmp[0]); ++ } ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp +new file mode 100644 +index 00000000..495f966b +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp +@@ -0,0 +1,20 @@ ++#version 450 ++ ++#include "generic_head.comp" ++#include "types.comp" ++ ++#extension GL_EXT_control_flow_attributes : enable ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++void main() { ++ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; ++ ++ if (i >= p.KX) { ++ return; ++ } ++ data_d[i] = D_TYPE(1. - 2. / (exp(2.*data_a[i]) + 1.)); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp +new file mode 100644 +index 00000000..28eb24e1 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp +@@ -0,0 +1,7 @@ ++#version 460 ++ ++#extension GL_NV_cooperative_matrix2 : require ++ ++void main() ++{ ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp +new file mode 100644 +index 00000000..79e065a9 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp +@@ -0,0 +1,41 @@ ++#version 450 ++ ++#extension GL_EXT_shader_16bit_storage : require ++ ++layout (push_constant) uniform parameter ++{ ++ uint nb1; ++ uint dim; ++ uint max_period; ++} p; ++ ++#include "types.comp" ++ ++#extension GL_EXT_control_flow_attributes : enable ++#define BLOCK_SIZE 256 ++ ++layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++void main() { ++ const uint i = gl_WorkGroupID.y; ++ const uint j = gl_GlobalInvocationID.x; ++ const uint d_offset = i * p.nb1; ++ ++ if (p.dim % 2 != 0 && j == ((p.dim + 1) / 2)) { ++ data_d[d_offset + p.dim] = 0.f; ++ } ++ ++ const uint half_dim = p.dim / 2; ++ if (j >= half_dim) { ++ return; ++ } ++ ++ const float timestep = float(data_a[i]); ++ const float freq = float(exp(-log(p.max_period) * j / half_dim)); ++ const float arg = timestep * freq; ++ data_d[d_offset + j] = D_TYPE(cos(arg)); ++ data_d[d_offset + j + half_dim] = D_TYPE(sin(arg)); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp +new file mode 100644 +index 00000000..eecc47f3 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp +@@ -0,0 +1,323 @@ ++ ++#if !defined(GGML_TYPES_COMP) ++#define GGML_TYPES_COMP ++ ++#extension GL_EXT_shader_explicit_arithmetic_types : require ++ ++#if defined(DATA_A_F32) ++#define QUANT_K 1 ++#define QUANT_R 1 ++ ++#if !defined(LOAD_VEC_A) || LOAD_VEC_A == 1 ++#define A_TYPE float ++#elif LOAD_VEC_A == 4 ++#define A_TYPE vec4 ++#elif LOAD_VEC_A == 8 ++#define A_TYPE mat2x4 ++#endif ++#endif ++ ++#if defined(DATA_A_F16) ++#define QUANT_K 1 ++#define QUANT_R 1 ++ ++#if !defined(LOAD_VEC_A) || LOAD_VEC_A == 1 ++#define A_TYPE float16_t ++#elif LOAD_VEC_A == 4 ++#define A_TYPE f16vec4 ++#elif LOAD_VEC_A == 8 ++#define A_TYPE f16mat2x4 ++#endif ++#endif ++ ++#define QUANT_K_Q4_0 32 ++#define QUANT_R_Q4_0 2 ++ ++struct block_q4_0 ++{ ++ float16_t d; ++ uint8_t qs[16]; ++}; ++struct block_q4_0_packed16 ++{ ++ float16_t d; ++ uint16_t qs[16/2]; ++}; ++ ++#if defined(DATA_A_Q4_0) ++#define QUANT_K QUANT_K_Q4_0 ++#define QUANT_R QUANT_R_Q4_0 ++#define A_TYPE block_q4_0 ++#define A_TYPE_PACKED16 block_q4_0_packed16 ++#endif ++ ++#define QUANT_K_Q4_1 32 ++#define QUANT_R_Q4_1 2 ++ ++struct block_q4_1 ++{ ++ float16_t d; ++ float16_t m; ++ uint8_t qs[16]; ++}; ++ ++struct block_q4_1_packed16 ++{ ++ float16_t d; ++ float16_t m; ++ uint16_t qs[16/2]; ++}; ++ ++#if defined(DATA_A_Q4_1) ++#define QUANT_K QUANT_K_Q4_1 ++#define QUANT_R QUANT_R_Q4_1 ++#define A_TYPE block_q4_1 ++#define A_TYPE_PACKED16 block_q4_1_packed16 ++#endif ++ ++#define QUANT_K_Q5_0 32 ++#define QUANT_R_Q5_0 2 ++ ++struct block_q5_0 ++{ ++ float16_t d; ++ uint16_t qh[2]; ++ uint8_t qs[16]; ++}; ++ ++struct block_q5_0_packed16 ++{ ++ float16_t d; ++ uint16_t qh[2]; ++ uint16_t qs[16/2]; ++}; ++ ++#if defined(DATA_A_Q5_0) ++#define QUANT_K QUANT_K_Q5_0 ++#define QUANT_R QUANT_R_Q5_0 ++#define A_TYPE block_q5_0 ++#define A_TYPE_PACKED16 block_q5_0_packed16 ++#endif ++ ++#define QUANT_K_Q5_1 32 ++#define QUANT_R_Q5_1 2 ++ ++struct block_q5_1 ++{ ++ float16_t d; ++ float16_t m; ++ uint qh; ++ uint8_t qs[16]; ++}; ++ ++struct block_q5_1_packed16 ++{ ++ float16_t d; ++ float16_t m; ++ uint qh; ++ uint16_t qs[16/2]; ++}; ++ ++#if defined(DATA_A_Q5_1) ++#define QUANT_K QUANT_K_Q5_1 ++#define QUANT_R QUANT_R_Q5_1 ++#define A_TYPE block_q5_1 ++#define A_TYPE_PACKED16 block_q5_1_packed16 ++#endif ++ ++#define QUANT_K_Q8_0 32 ++#define QUANT_R_Q8_0 1 ++ ++struct block_q8_0 ++{ ++ float16_t d; ++ int8_t qs[32]; ++}; ++struct block_q8_0_packed16 ++{ ++ float16_t d; ++ uint16_t qs[32/2]; ++}; ++ ++#if defined(DATA_A_Q8_0) ++#define QUANT_K QUANT_K_Q8_0 ++#define QUANT_R QUANT_R_Q8_0 ++#define A_TYPE block_q8_0 ++#define A_TYPE_PACKED16 block_q8_0_packed16 ++#endif ++ ++// K-quants ++#define QUANT_K_Q2_K 256 ++ ++struct block_q2_K ++{ ++ uint8_t scales[QUANT_K_Q2_K/16]; ++ uint8_t qs[QUANT_K_Q2_K/4]; ++ f16vec2 d; ++}; ++ ++struct block_q2_K_packed16 ++{ ++ uint16_t scales[QUANT_K_Q2_K/16/2]; ++ uint16_t qs[QUANT_K_Q2_K/4/2]; ++ f16vec2 d; ++}; ++ ++struct block_q2_K_packed32 ++{ ++ uint32_t scales[QUANT_K_Q2_K/16/4]; ++ uint32_t qs[QUANT_K_Q2_K/4/4]; ++ f16vec2 d; ++}; ++ ++#if defined(DATA_A_Q2_K) ++#define QUANT_K QUANT_K_Q2_K ++#define A_TYPE block_q2_K ++#define A_TYPE_PACKED16 block_q2_K_packed16 ++#define A_TYPE_PACKED32 block_q2_K_packed32 ++#endif ++ ++#define QUANT_K_Q3_K 256 ++ ++struct block_q3_K ++{ ++ uint8_t hmask[QUANT_K_Q3_K/8]; ++ uint8_t qs[QUANT_K_Q3_K/4]; ++ uint8_t scales[12]; ++ float16_t d; ++}; ++ ++struct block_q3_K_packed16 ++{ ++ uint16_t hmask[QUANT_K_Q3_K/8/2]; ++ uint16_t qs[QUANT_K_Q3_K/4/2]; ++ uint16_t scales[12/2]; ++ float16_t d; ++}; ++ ++#if defined(DATA_A_Q3_K) ++#define QUANT_K QUANT_K_Q3_K ++#define A_TYPE block_q3_K ++#define A_TYPE_PACKED16 block_q3_K_packed16 ++#endif ++ ++#define QUANT_K_Q4_K 256 ++ ++struct block_q4_K ++{ ++ f16vec2 d; ++ uint8_t scales[3*QUANT_K_Q4_K/64]; ++ uint8_t qs[QUANT_K_Q4_K/2]; ++}; ++ ++struct block_q4_K_packed16 ++{ ++ f16vec2 d; ++ uint16_t scales[3*QUANT_K_Q4_K/64/2]; ++ uint16_t qs[QUANT_K_Q4_K/2/2]; ++}; ++ ++struct block_q4_K_packed32 ++{ ++ f16vec2 d; ++ uint32_t scales[3*QUANT_K_Q4_K/64/4]; ++ uint32_t qs[QUANT_K_Q4_K/2/4]; ++}; ++ ++#if defined(DATA_A_Q4_K) ++#define QUANT_K QUANT_K_Q4_K ++#define A_TYPE block_q4_K ++#define A_TYPE_PACKED16 block_q4_K_packed16 ++#define A_TYPE_PACKED32 block_q4_K_packed32 ++#endif ++ ++#define QUANT_K_Q5_K 256 ++ ++struct block_q5_K ++{ ++ f16vec2 d; ++ uint8_t scales[12]; ++ uint8_t qh[QUANT_K_Q5_K/8]; ++ uint8_t qs[QUANT_K_Q5_K/2]; ++}; ++ ++struct block_q5_K_packed16 ++{ ++ f16vec2 d; ++ uint16_t scales[12/2]; ++ uint16_t qh[QUANT_K_Q5_K/8/2]; ++ uint16_t qs[QUANT_K_Q5_K/2/2]; ++}; ++ ++#if defined(DATA_A_Q5_K) ++#define QUANT_K QUANT_K_Q5_K ++#define A_TYPE block_q5_K ++#define A_TYPE_PACKED16 block_q5_K_packed16 ++#endif ++ ++#define QUANT_K_Q6_K 256 ++ ++struct block_q6_K ++{ ++ uint8_t ql[QUANT_K_Q6_K/2]; ++ uint8_t qh[QUANT_K_Q6_K/4]; ++ int8_t scales[QUANT_K_Q6_K/16]; ++ float16_t d; ++}; ++ ++struct block_q6_K_packed16 ++{ ++ uint16_t ql[QUANT_K_Q6_K/2/2]; ++ uint16_t qh[QUANT_K_Q6_K/4/2]; ++ int8_t scales[QUANT_K_Q6_K/16]; ++ float16_t d; ++}; ++ ++#if defined(DATA_A_Q6_K) ++#define QUANT_K QUANT_K_Q6_K ++#define A_TYPE block_q6_K ++#define A_TYPE_PACKED16 block_q6_K_packed16 ++#endif ++ ++// IQuants ++ ++#define QUANT_K_IQ4_NL 32 ++#define QUANT_R_IQ4_NL 2 ++ ++struct block_iq4_nl ++{ ++ float16_t d; ++ uint8_t qs[QUANT_K_IQ4_NL/2]; ++}; ++ ++struct block_iq4_nl_packed16 ++{ ++ float16_t d; ++ uint16_t qs[QUANT_K_IQ4_NL/2/2]; ++}; ++ ++#if defined(DATA_A_IQ4_NL) ++ ++const int8_t kvalues_iq4nl_const[16] = { ++ int8_t(-127), int8_t(-104), int8_t(-83), int8_t(-65), int8_t(-49), int8_t(-35), int8_t(-22), int8_t(-10), ++ int8_t(1), int8_t(13), int8_t(25), int8_t(38), int8_t(53), int8_t(69), int8_t(89), int8_t(113) ++}; ++ ++shared FLOAT_TYPE kvalues_iq4nl[16]; ++ ++void init_iq4nl_shmem() ++{ ++ // copy the table into shared memory and sync ++ if (gl_LocalInvocationIndex.x < 16) { ++ kvalues_iq4nl[gl_LocalInvocationIndex.x] = FLOAT_TYPE(kvalues_iq4nl_const[gl_LocalInvocationIndex.x]); ++ } ++ barrier(); ++} ++ ++#define QUANT_K QUANT_K_IQ4_NL ++#define QUANT_R QUANT_R_IQ4_NL ++#define A_TYPE block_iq4_nl ++#define A_TYPE_PACKED16 block_iq4_nl_packed16 ++#endif ++ ++#endif // !defined(GGML_TYPES_COMP) +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp +new file mode 100644 +index 00000000..6f607380 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp +@@ -0,0 +1,36 @@ ++#version 450 ++ ++layout (push_constant) uniform parameter ++{ ++ uint ne; uint a_offset; uint d_offset; ++ uint nb00; uint nb01; uint nb02; uint nb03; ++ uint ne10; uint ne11; uint ne12; uint ne13; ++ float sf0; float sf1; float sf2; float sf3; ++} p; ++ ++#include "types.comp" ++ ++layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; ++ ++layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; ++layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; ++ ++void main() { ++ const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; ++ ++ if (idx >= p.ne) { ++ return; ++ } ++ ++ const uint i10 = idx % p.ne10; ++ const uint i11 = (idx / p.ne10) % p.ne11; ++ const uint i12 = (idx / (p.ne10 * p.ne11)) % p.ne12; ++ const uint i13 = (idx / (p.ne10 * p.ne11 * p.ne12)) % p.ne13; ++ ++ const uint i00 = uint(i10 / p.sf0); ++ const uint i01 = uint(i11 / p.sf1); ++ const uint i02 = uint(i12 / p.sf2); ++ const uint i03 = uint(i13 / p.sf3); ++ ++ data_d[p.d_offset + idx] = D_TYPE(data_a[p.a_offset + i03 * p.nb03 + i02 * p.nb02 + i01 * p.nb01 + i00 * p.nb00]); ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +new file mode 100644 +index 00000000..8111c063 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +@@ -0,0 +1,594 @@ ++ ++ ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++ ++#ifdef _WIN32 ++ #include ++ #include // For _mkdir on Windows ++ #include // For std::replace on w64devkit ++#else ++ #include ++ #include ++ #include ++#endif ++ ++#include ++ ++#define ASYNCIO_CONCURRENCY 64 ++ ++std::mutex lock; ++std::vector> shader_fnames; ++ ++std::string GLSLC = "glslc"; ++std::string input_dir = "vulkan-shaders"; ++std::string output_dir = "/tmp"; ++std::string target_hpp = "ggml-vulkan-shaders.hpp"; ++std::string target_cpp = "ggml-vulkan-shaders.cpp"; ++bool no_clean = false; ++ ++const std::vector type_names = { ++ "f32", ++ "f16", ++ "q4_0", ++ "q4_1", ++ "q5_0", ++ "q5_1", ++ "q8_0", ++ "q2_k", ++ "q3_k", ++ "q4_k", ++ "q5_k", ++ "q6_k", ++ "iq4_nl" ++}; ++ ++namespace { ++void execute_command(const std::string& command, std::string& stdout_str, std::string& stderr_str) { ++#ifdef _WIN32 ++ HANDLE stdout_read, stdout_write; ++ HANDLE stderr_read, stderr_write; ++ SECURITY_ATTRIBUTES sa = { sizeof(SECURITY_ATTRIBUTES), NULL, TRUE }; ++ ++ if (!CreatePipe(&stdout_read, &stdout_write, &sa, 0) || ++ !SetHandleInformation(stdout_read, HANDLE_FLAG_INHERIT, 0)) { ++ throw std::runtime_error("Failed to create stdout pipe"); ++ } ++ ++ if (!CreatePipe(&stderr_read, &stderr_write, &sa, 0) || ++ !SetHandleInformation(stderr_read, HANDLE_FLAG_INHERIT, 0)) { ++ throw std::runtime_error("Failed to create stderr pipe"); ++ } ++ ++ PROCESS_INFORMATION pi; ++ STARTUPINFOA si = {}; ++ si.cb = sizeof(STARTUPINFOA); ++ si.dwFlags = STARTF_USESTDHANDLES; ++ si.hStdOutput = stdout_write; ++ si.hStdError = stderr_write; ++ ++ std::vector cmd(command.begin(), command.end()); ++ cmd.push_back('\0'); ++ ++ if (!CreateProcessA(NULL, cmd.data(), NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi)) { ++ throw std::runtime_error("Failed to create process"); ++ } ++ ++ CloseHandle(stdout_write); ++ CloseHandle(stderr_write); ++ ++ std::array buffer; ++ DWORD bytes_read; ++ ++ while (ReadFile(stdout_read, buffer.data(), (DWORD)buffer.size(), &bytes_read, NULL) && bytes_read > 0) { ++ stdout_str.append(buffer.data(), bytes_read); ++ } ++ ++ while (ReadFile(stderr_read, buffer.data(), (DWORD)buffer.size(), &bytes_read, NULL) && bytes_read > 0) { ++ stderr_str.append(buffer.data(), bytes_read); ++ } ++ ++ CloseHandle(stdout_read); ++ CloseHandle(stderr_read); ++ WaitForSingleObject(pi.hProcess, INFINITE); ++ CloseHandle(pi.hProcess); ++ CloseHandle(pi.hThread); ++#else ++int stdout_pipe[2]; ++ int stderr_pipe[2]; ++ ++ if (pipe(stdout_pipe) != 0 || pipe(stderr_pipe) != 0) { ++ throw std::runtime_error("Failed to create pipes"); ++ } ++ ++ pid_t pid = fork(); ++ if (pid < 0) { ++ throw std::runtime_error("Failed to fork process"); ++ } ++ ++ if (pid == 0) { ++ close(stdout_pipe[0]); ++ close(stderr_pipe[0]); ++ dup2(stdout_pipe[1], STDOUT_FILENO); ++ dup2(stderr_pipe[1], STDERR_FILENO); ++ close(stdout_pipe[1]); ++ close(stderr_pipe[1]); ++ execl("/bin/sh", "sh", "-c", command.c_str(), (char*) nullptr); ++ _exit(EXIT_FAILURE); ++ } else { ++ close(stdout_pipe[1]); ++ close(stderr_pipe[1]); ++ ++ std::array buffer; ++ ssize_t bytes_read; ++ ++ while ((bytes_read = read(stdout_pipe[0], buffer.data(), buffer.size())) > 0) { ++ stdout_str.append(buffer.data(), bytes_read); ++ } ++ ++ while ((bytes_read = read(stderr_pipe[0], buffer.data(), buffer.size())) > 0) { ++ stderr_str.append(buffer.data(), bytes_read); ++ } ++ ++ close(stdout_pipe[0]); ++ close(stderr_pipe[0]); ++ waitpid(pid, nullptr, 0); ++ } ++#endif ++} ++ ++bool directory_exists(const std::string& path) { ++ struct stat info; ++ if (stat(path.c_str(), &info) != 0) { ++ return false; // Path doesn't exist or can't be accessed ++ } ++ return (info.st_mode & S_IFDIR) != 0; // Check if it is a directory ++} ++ ++bool create_directory(const std::string& path) { ++#ifdef _WIN32 ++ return _mkdir(path.c_str()) == 0 || errno == EEXIST; // EEXIST means the directory already exists ++#else ++ return mkdir(path.c_str(), 0755) == 0 || errno == EEXIST; // 0755 is the directory permissions ++#endif ++} ++ ++std::string to_uppercase(const std::string& input) { ++ std::string result = input; ++ for (char& c : result) { ++ c = std::toupper(c); ++ } ++ return result; ++} ++ ++bool string_ends_with(const std::string& str, const std::string& suffix) { ++ if (suffix.size() > str.size()) { ++ return false; ++ } ++ return std::equal(suffix.rbegin(), suffix.rend(), str.rbegin()); ++} ++ ++static const char path_separator = '/'; ++ ++std::string join_paths(const std::string& path1, const std::string& path2) { ++ return path1 + path_separator + path2; ++} ++ ++std::string basename(const std::string &path) { ++ return path.substr(path.find_last_of("/\\") + 1); ++} ++ ++// variables to track number of compiles in progress ++static uint32_t compile_count = 0; ++static std::mutex compile_count_mutex; ++static std::condition_variable compile_count_cond; ++ ++void string_to_spv_func(const std::string& _name, const std::string& in_fname, const std::map& defines, bool fp16 = true, bool coopmat = false, bool coopmat2 = false, bool f16acc = false) { ++ std::string name = _name + (f16acc ? "_f16acc" : "") + (coopmat ? "_coopmat" : "") + (coopmat2 ? "_cm2" : (fp16 ? "" : "_fp32")); ++ std::string out_fname = join_paths(output_dir, name + ".spv"); ++ std::string in_path = join_paths(input_dir, in_fname); ++ ++ std::string target_env = (name.find("_cm2") != std::string::npos) ? "--target-env=vulkan1.3" : "--target-env=vulkan1.2"; ++ ++ // disable spirv-opt for coopmat shaders for https://github.com/ggerganov/llama.cpp/issues/10734 ++ std::string opt_level = coopmat ? "" : "-O"; ++ ++ #ifdef _WIN32 ++ std::vector cmd = {GLSLC, "-fshader-stage=compute", target_env, opt_level, "\"" + in_path + "\"", "-o", "\"" + out_fname + "\""}; ++ #else ++ std::vector cmd = {GLSLC, "-fshader-stage=compute", target_env, opt_level, in_path, "-o", out_fname}; ++ #endif ++ ++ #ifdef GGML_VULKAN_SHADER_DEBUG_INFO ++ cmd.push_back("-g"); ++ #endif ++ ++ for (const auto& define : defines) { ++ cmd.push_back("-D" + define.first + "=" + define.second); ++ } ++ ++ std::string command; ++ for (const auto& part : cmd) { ++ command += part + " "; ++ } ++ ++ std::string stdout_str, stderr_str; ++ try { ++ // std::cout << "Executing command: "; ++ // for (const auto& part : cmd) { ++ // std::cout << part << " "; ++ // } ++ // std::cout << std::endl; ++ ++ execute_command(command, stdout_str, stderr_str); ++ if (!stderr_str.empty()) { ++ std::cerr << "cannot compile " << name << "\n\n" << command << "\n\n" << stderr_str << std::endl; ++ return; ++ } ++ ++ std::lock_guard guard(lock); ++ shader_fnames.push_back(std::make_pair(name, out_fname)); ++ } catch (const std::exception& e) { ++ std::cerr << "Error executing command for " << name << ": " << e.what() << std::endl; ++ } ++ { ++ std::lock_guard guard(compile_count_mutex); ++ assert(compile_count > 0); ++ compile_count--; ++ } ++ compile_count_cond.notify_all(); ++} ++ ++std::map merge_maps(const std::map& a, const std::map& b) { ++ std::map result = a; ++ result.insert(b.begin(), b.end()); ++ return result; ++} ++ ++static std::vector> compiles; ++void string_to_spv(const std::string& _name, const std::string& in_fname, const std::map& defines, bool fp16 = true, bool coopmat = false, bool coopmat2 = false, bool f16acc = false) { ++ { ++ // wait until fewer than N compiles are in progress. ++ // 16 is an arbitrary limit, the goal is to avoid "failed to create pipe" errors. ++ uint32_t N = 16; ++ std::unique_lock guard(compile_count_mutex); ++ while (compile_count >= N) { ++ compile_count_cond.wait(guard); ++ } ++ compile_count++; ++ } ++ compiles.push_back(std::async(string_to_spv_func, _name, in_fname, defines, fp16, coopmat, coopmat2, f16acc)); ++} ++ ++void matmul_shaders(bool fp16, bool matmul_id, bool coopmat, bool coopmat2, bool f16acc) { ++ std::string load_vec = coopmat2 ? "1" : fp16 ? "8" : "4"; ++ std::string aligned_b_type_f32 = coopmat2 ? "float" : fp16 ? "mat2x4" : "vec4"; ++ std::string aligned_b_type_f16 = coopmat2 ? "float16_t" : fp16 ? "f16mat2x4" : "f16vec4"; ++ ++ std::map base_dict = {{"FLOAT_TYPE", (coopmat2 || fp16) ? "float16_t" : "float"}}; ++ std::string shader_name = "matmul"; ++ ++ if (matmul_id) { ++ base_dict["MUL_MAT_ID"] = "1"; ++ shader_name = "matmul_id"; ++ } ++ ++ if (fp16) { ++ base_dict["FLOAT16"] = "1"; ++ } ++ ++ base_dict["ACC_TYPE"] = f16acc ? "float16_t" : "float"; ++ ++ if (coopmat) { ++ base_dict["COOPMAT"] = "1"; ++ } ++ ++ base_dict["ACC_TYPE"] = f16acc ? "float16_t" : "float"; ++ ++ std::string source_name = coopmat2 ? "mul_mm_cm2.comp" : "mul_mm.comp"; ++ ++ // Shaders with f16 B_TYPE ++ string_to_spv(shader_name + "_f32_f16", source_name, merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}, }), fp16, coopmat, coopmat2, f16acc); ++ string_to_spv(shader_name + "_f32_f16_aligned", source_name, merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); ++ ++ string_to_spv(shader_name + "_f16_aligned", source_name, merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); ++ string_to_spv(shader_name + "_f16", source_name, merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); ++ ++ for (const auto& tname : type_names) { ++ std::string data_a_key = "DATA_A_" + to_uppercase(tname); ++ // For unaligned, load one at a time for f32/f16, or two at a time for quants ++ std::string load_vec_a_unaligned = (coopmat2 || tname == "f32" || tname == "f16") ? "1" : "2"; ++ // For aligned matmul loads ++ std::string load_vec_a = (coopmat2 || tname == "f32" || tname == "f16") ? load_vec : "2"; ++ ++ string_to_spv(shader_name + "_" + tname + "_f32", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}}), fp16, coopmat, coopmat2, f16acc); ++ string_to_spv(shader_name + "_" + tname + "_f32_aligned", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); ++ ++ if (tname != "f16" && tname != "f32") { ++ string_to_spv(shader_name + "_" + tname + "_f16", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}}), fp16, coopmat, coopmat2, f16acc); ++ string_to_spv(shader_name + "_" + tname + "_f16_aligned", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); ++ } ++ } ++} ++ ++void process_shaders() { ++ std::cout << "ggml_vulkan: Generating and compiling shaders to SPIR-V" << std::endl; ++ std::map base_dict = {{"FLOAT_TYPE", "float"}}; ++ ++ // matmul ++ for (const auto& matmul_id : {false, true}) { ++ // No coopmats ++ // fp32 ++ matmul_shaders(false, matmul_id, false, false, false); ++ ++ // fp16, fp32acc and fp16acc ++ matmul_shaders(true, matmul_id, false, false, false); ++ matmul_shaders(true, matmul_id, false, false, true); ++ ++ // Coopmat, fp32acc and fp16acc ++ matmul_shaders(true, matmul_id, true, false, false); ++ matmul_shaders(true, matmul_id, true, false, true); ++ ++#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) ++ // Coopmat2, fp32acc and fp16acc ++ matmul_shaders(true, matmul_id, false, true, false); ++ matmul_shaders(true, matmul_id, false, true, true); ++#endif ++ } ++ ++#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) ++ // flash attention ++ for (const auto& f16acc : {false, true}) { ++ std::string acctype = f16acc ? "float16_t" : "float"; ++ ++ for (const auto& tname : type_names) { ++ if (tname == "f32") { ++ continue; ++ } ++ ++ if (tname == "f16") { ++ string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm2.comp", ++ merge_maps(base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"ACC_TYPE", acctype}}), true, false, true, f16acc); ++ } else { ++ std::string data_a_key = "DATA_A_" + to_uppercase(tname); ++ string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm2.comp", ++ merge_maps(base_dict, {{data_a_key, "1"}, {"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"ACC_TYPE", acctype}, {"DEQUANTFUNC", "dequantFunc"+to_uppercase(tname) }, {"BLOCK_SIZE", "QUANT_K_"+to_uppercase(tname) }}), true, false, true, f16acc); ++ } ++ } ++ } ++#endif ++ ++ for (const auto& tname : type_names) { ++ // mul mat vec ++ std::string data_a_key = "DATA_A_" + to_uppercase(tname); ++ std::string shader = (string_ends_with(tname, "_k")) ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp"; ++ ++ string_to_spv("mul_mat_vec_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPE_VEC2", "vec2"}, {"B_TYPE_VEC4", "vec4"}, {"D_TYPE", "float"}})); ++ string_to_spv("mul_mat_vec_" + tname + "_f16_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"B_TYPE_VEC2", "f16vec2"}, {"B_TYPE_VEC4", "f16vec4"}, {"D_TYPE", "float"}})); ++ ++ string_to_spv("mul_mat_vec_id_" + tname + "_f32", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPE_VEC2", "vec2"}, {"B_TYPE_VEC4", "vec4"}, {"D_TYPE", "float"}})); ++ ++ // Dequant shaders ++ if (tname != "f16") { ++ string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}})); ++ } ++ ++ if (!string_ends_with(tname, "_k")) { ++ shader = (tname == "f32" || tname == "f16") ? "get_rows.comp" : "get_rows_quant.comp"; ++ ++ if (tname == "f16") { ++ string_to_spv("get_rows_" + tname, shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}})); ++ } else { ++ string_to_spv("get_rows_" + tname, shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}})); ++ } ++ string_to_spv("get_rows_" + tname + "_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float"}})); ++ } ++ } ++ ++ string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); ++ string_to_spv("mul_mat_vec_nc_f16_f32", "mul_mat_vec_nc.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); ++ ++ // Norms ++ string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); ++ string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); ++ string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); ++ ++ string_to_spv("cpy_f32_f32", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ string_to_spv("cpy_f32_f16", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); ++ string_to_spv("cpy_f16_f16", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); ++ string_to_spv("contig_cpy_f32_f32", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ string_to_spv("contig_cpy_f32_f16", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); ++ string_to_spv("contig_cpy_f16_f16", "contig_copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); ++ ++ string_to_spv("add_f32", "add.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); ++ string_to_spv("add_f16_f32_f16", "add.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"FLOAT_TYPE", "float"}}); ++ ++ string_to_spv("acc_f32", "acc.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); ++ ++ string_to_spv("split_k_reduce", "mul_mat_split_k_reduce.comp", {}); ++ ++ string_to_spv("mul_f32", "mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); ++ ++ string_to_spv("div_f32", "div.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); ++ ++ string_to_spv("repeat_f32", "repeat.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ ++ string_to_spv("scale_f32", "scale.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); ++ ++ string_to_spv("sqr_f32", "square.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); ++ ++ string_to_spv("sin_f32", "sin.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); ++ ++ string_to_spv("cos_f32", "cos.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); ++ ++ string_to_spv("clamp_f32", "clamp.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); ++ ++ string_to_spv("pad_f32", "pad.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ ++ string_to_spv("concat_f32", "concat.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); ++ string_to_spv("concat_f16", "concat.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); ++ string_to_spv("concat_i32", "concat.comp", {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}}); ++ ++ string_to_spv("upscale_f32", "upscale.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); ++ ++ string_to_spv("gelu_f32", "gelu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ string_to_spv("gelu_quick_f32", "gelu_quick.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ string_to_spv("silu_f32", "silu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ string_to_spv("relu_f32", "relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ string_to_spv("leaky_relu_f32", "leaky_relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ string_to_spv("tanh_f32", "tanh.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ ++ string_to_spv("diag_mask_inf_f32", "diag_mask_inf.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ ++ string_to_spv("soft_max_f32", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); ++ string_to_spv("soft_max_f32_f16", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}})); ++ ++ string_to_spv("rope_norm_f32", "rope_norm.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ string_to_spv("rope_norm_f16", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); ++ string_to_spv("rope_norm_f16_rte", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}}); ++ ++ string_to_spv("rope_neox_f32", "rope_neox.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); ++ string_to_spv("rope_neox_f16", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); ++ string_to_spv("rope_neox_f16_rte", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}}); ++ ++ string_to_spv("argsort_f32", "argsort.comp", {{"A_TYPE", "float"}}); ++ ++ string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); ++ ++ string_to_spv("im2col_f32", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); ++ string_to_spv("im2col_f32_f16", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}})); ++ string_to_spv("im2col_f32_f16_rte", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}})); ++ ++ string_to_spv("timestep_embedding_f32", "timestep_embedding.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); ++ ++ string_to_spv("pool2d_f32", "pool2d.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); ++ ++ string_to_spv("rwkv_wkv6_f32", "wkv6.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); ++ ++ for (auto &c : compiles) { ++ c.wait(); ++ } ++} ++ ++void write_output_files() { ++ FILE* hdr = fopen(target_hpp.c_str(), "w"); ++ FILE* src = fopen(target_cpp.c_str(), "w"); ++ ++ fprintf(hdr, "#include \n\n"); ++ fprintf(src, "#include \"%s\"\n\n", basename(target_hpp).c_str()); ++ ++ for (const auto& pair : shader_fnames) { ++ const std::string& name = pair.first; ++ #ifdef _WIN32 ++ std::string path = pair.second; ++ std::replace(path.begin(), path.end(), '/', '\\' ); ++ #else ++ const std::string& path = pair.second; ++ #endif ++ ++ FILE* spv = fopen(path.c_str(), "rb"); ++ if (!spv) { ++ std::cerr << "Error opening SPIR-V file: " << path << " (" << strerror(errno) << ")\n"; ++ continue; ++ } ++ ++ fseek(spv, 0, SEEK_END); ++ size_t size = ftell(spv); ++ fseek(spv, 0, SEEK_SET); ++ ++ std::vector data(size); ++ size_t read_size = fread(data.data(), 1, size, spv); ++ fclose(spv); ++ if (read_size != size) { ++ std::cerr << "Error reading SPIR-V file: " << path << " (" << strerror(errno) << ")\n"; ++ continue; ++ } ++ ++ fprintf(hdr, "extern unsigned char %s_data[%zu];\n", name.c_str(), size); ++ fprintf(hdr, "const uint64_t %s_len = %zu;\n\n", name.c_str(), size); ++ ++ fprintf(src, "unsigned char %s_data[%zu] = {\n", name.c_str(), size); ++ for (size_t i = 0; i < size; ++i) { ++ fprintf(src, "0x%02x,", data[i]); ++ if ((i + 1) % 12 == 0) fprintf(src, "\n"); ++ } ++ fprintf(src, "\n};\n\n"); ++ ++ if (!no_clean) { ++ std::remove(path.c_str()); ++ } ++ } ++ ++ fclose(hdr); ++ fclose(src); ++} ++} ++ ++int main(int argc, char** argv) { ++ std::map args; ++ for (int i = 1; i < argc; ++i) { ++ std::string arg = argv[i]; ++ if (arg.rfind("--", 0) == 0) { ++ if (i + 1 < argc && argv[i + 1][0] != '-') { ++ args[arg] = argv[i + 1]; ++ ++i; ++ } else { ++ args[arg] = ""; ++ } ++ } ++ } ++ ++ if (args.find("--glslc") != args.end()) { ++ GLSLC = args["--glslc"]; // Path to glslc ++ } ++ if (args.find("--input-dir") != args.end()) { ++ input_dir = args["--input-dir"]; // Directory containing shader sources ++ } ++ if (args.find("--output-dir") != args.end()) { ++ output_dir = args["--output-dir"]; // Directory for containing SPIR-V output ++ } ++ if (args.find("--target-hpp") != args.end()) { ++ target_hpp = args["--target-hpp"]; // Path to generated header file ++ } ++ if (args.find("--target-cpp") != args.end()) { ++ target_cpp = args["--target-cpp"]; // Path to generated cpp file ++ } ++ if (args.find("--no-clean") != args.end()) { ++ no_clean = true; // Keep temporary SPIR-V files in output-dir after build ++ } ++ ++ if (!directory_exists(input_dir)) { ++ std::cerr << "\"" << input_dir << "\" must be a valid directory containing shader sources" << std::endl; ++ return EXIT_FAILURE; ++ } ++ ++ if (!directory_exists(output_dir)) { ++ if (!create_directory(output_dir)) { ++ std::cerr << "Error creating output directory: " << output_dir << "\n"; ++ return EXIT_FAILURE; ++ } ++ } ++ ++ process_shaders(); ++ ++ write_output_files(); ++ ++ return EXIT_SUCCESS; ++} +diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp +new file mode 100644 +index 00000000..35cc6c45 +--- /dev/null ++++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp +@@ -0,0 +1,87 @@ ++#version 450 ++ ++#extension GL_EXT_control_flow_attributes : require ++ ++#define BLOCK_SIZE 64 ++layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; ++ ++layout(push_constant) uniform Parameters { ++ uint B; ++ uint T; ++ uint C; ++ uint H; ++}; ++ ++layout(binding = 0) readonly buffer KBuf { A_TYPE k[]; }; ++layout(binding = 1) readonly buffer VBuf { A_TYPE v[]; }; ++layout(binding = 2) readonly buffer RBuf { A_TYPE r[]; }; ++layout(binding = 3) readonly buffer TimeFBuf { A_TYPE tf[]; }; ++layout(binding = 4) readonly buffer TimeDBuf { A_TYPE td[]; }; ++layout(binding = 5) readonly buffer StateBuf { A_TYPE state_in[]; }; ++layout(binding = 6) buffer DstBuf { A_TYPE dst[]; }; ++ ++shared A_TYPE _k[BLOCK_SIZE], _r[BLOCK_SIZE], _tf[BLOCK_SIZE], _td[BLOCK_SIZE]; ++ ++void main() { ++ const uint head_size = BLOCK_SIZE; ++ const uint batch_id = gl_WorkGroupID.x / H; ++ const uint head_id = gl_WorkGroupID.x % H; ++ const uint tid = gl_LocalInvocationID.x; ++ ++ const uint state_size = C * head_size; ++ const uint n_seq_tokens = T / B; ++ ++ if (batch_id >= B || head_id >= H) { ++ return; ++ } ++ ++ A_TYPE state[BLOCK_SIZE]; ++ [[unroll]] for (uint i = 0; i < head_size; i++) { ++ state[i] = state_in[batch_id * state_size + head_id * head_size * head_size ++ + i * head_size + tid]; ++ } ++ ++ barrier(); ++ _tf[tid] = tf[head_id * head_size + tid]; ++ barrier(); ++ ++ const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid; ++ const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid; ++ ++ for (uint t = start_t; t < end_t; t += C) { ++ barrier(); ++ _k[tid] = k[t]; ++ _r[tid] = r[t]; ++ _td[tid] = td[t]; ++ barrier(); ++ ++ const A_TYPE v_val = v[t]; ++ A_TYPE y = 0.0; ++ ++ [[unroll]] for (uint j = 0; j < head_size; j += 4) { ++ vec4 k_vec = vec4(_k[j], _k[j+1], _k[j+2], _k[j+3]); ++ vec4 r_vec = vec4(_r[j], _r[j+1], _r[j+2], _r[j+3]); ++ vec4 tf_vec = vec4(_tf[j], _tf[j+1], _tf[j+2], _tf[j+3]); ++ vec4 td_vec = vec4(_td[j], _td[j+1], _td[j+2], _td[j+3]); ++ vec4 s_vec = vec4(state[j], state[j+1], state[j+2], state[j+3]); ++ ++ vec4 kv = k_vec * v_val; ++ ++ vec4 temp = tf_vec * kv + s_vec; ++ y += dot(r_vec, temp); ++ ++ s_vec = s_vec * td_vec + kv; ++ state[j] = s_vec.x; ++ state[j+1] = s_vec.y; ++ state[j+2] = s_vec.z; ++ state[j+3] = s_vec.w; ++ } ++ ++ dst[t] = y; ++ } ++ ++ [[unroll]] for (uint i = 0; i < head_size; i++) { ++ dst[T * C + batch_id * state_size + head_id * head_size * head_size ++ + i * head_size + tid] = state[i]; ++ } ++} +-- +2.43.0 + From 98f699773aad00ba60c0cda4ca1a0ce7940d0a09 Mon Sep 17 00:00:00 2001 From: Vadim Grinco Date: Mon, 10 Mar 2025 12:34:37 +0100 Subject: [PATCH 25/29] Applied 00-fix-vulkan-building.patch Work done by McBane87 here: https://github.com/whyvl/ollama-vulkan/issues/7#issuecomment-2660836871 Signed-off-by: Vadim Grinco --- CMakePresets.json | 13 +- discover/gpu.go | 7 +- .../ggml/ggml/src/ggml-vulkan/CMakeLists.txt | 92 + .../ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp | 8745 +++++++++++++++++ .../ggml-vulkan/vulkan-shaders/CMakeLists.txt | 9 + .../src/ggml-vulkan/vulkan-shaders/acc.comp | 29 + .../src/ggml-vulkan/vulkan-shaders/add.comp | 29 + .../ggml-vulkan/vulkan-shaders/argsort.comp | 69 + .../src/ggml-vulkan/vulkan-shaders/clamp.comp | 17 + .../ggml-vulkan/vulkan-shaders/concat.comp | 41 + .../vulkan-shaders/contig_copy.comp | 42 + .../src/ggml-vulkan/vulkan-shaders/copy.comp | 20 + .../src/ggml-vulkan/vulkan-shaders/cos.comp | 17 + .../vulkan-shaders/dequant_f32.comp | 20 + .../vulkan-shaders/dequant_funcs.comp | 118 + .../vulkan-shaders/dequant_funcs_cm2.comp | 325 + .../vulkan-shaders/dequant_head.comp | 13 + .../vulkan-shaders/dequant_iq4_nl.comp | 32 + .../vulkan-shaders/dequant_q2_k.comp | 34 + .../vulkan-shaders/dequant_q3_k.comp | 42 + .../vulkan-shaders/dequant_q4_0.comp | 30 + .../vulkan-shaders/dequant_q4_1.comp | 32 + .../vulkan-shaders/dequant_q4_k.comp | 68 + .../vulkan-shaders/dequant_q5_0.comp | 34 + .../vulkan-shaders/dequant_q5_1.comp | 35 + .../vulkan-shaders/dequant_q5_k.comp | 70 + .../vulkan-shaders/dequant_q6_k.comp | 33 + .../vulkan-shaders/dequant_q8_0.comp | 31 + .../vulkan-shaders/diag_mask_inf.comp | 34 + .../src/ggml-vulkan/vulkan-shaders/div.comp | 27 + .../vulkan-shaders/flash_attn_cm2.comp | 289 + .../src/ggml-vulkan/vulkan-shaders/gelu.comp | 25 + .../vulkan-shaders/gelu_quick.comp | 23 + .../vulkan-shaders/generic_binary_head.comp | 64 + .../vulkan-shaders/generic_head.comp | 9 + .../vulkan-shaders/generic_unary_head.comp | 56 + .../ggml-vulkan/vulkan-shaders/get_rows.comp | 28 + .../vulkan-shaders/get_rows_quant.comp | 39 + .../vulkan-shaders/group_norm.comp | 66 + .../ggml-vulkan/vulkan-shaders/im2col.comp | 87 + .../vulkan-shaders/leaky_relu.comp | 22 + .../src/ggml-vulkan/vulkan-shaders/mul.comp | 27 + .../mul_mat_split_k_reduce.comp | 48 + .../vulkan-shaders/mul_mat_vec.comp | 152 + .../vulkan-shaders/mul_mat_vec_base.comp | 118 + .../vulkan-shaders/mul_mat_vec_nc.comp | 71 + .../vulkan-shaders/mul_mat_vec_p021.comp | 73 + .../vulkan-shaders/mul_mat_vec_q2_k.comp | 115 + .../vulkan-shaders/mul_mat_vec_q3_k.comp | 103 + .../vulkan-shaders/mul_mat_vec_q4_k.comp | 133 + .../vulkan-shaders/mul_mat_vec_q5_k.comp | 162 + .../vulkan-shaders/mul_mat_vec_q6_k.comp | 112 + .../ggml-vulkan/vulkan-shaders/mul_mm.comp | 631 ++ .../vulkan-shaders/mul_mm_cm2.comp | 328 + .../src/ggml-vulkan/vulkan-shaders/norm.comp | 44 + .../src/ggml-vulkan/vulkan-shaders/pad.comp | 28 + .../ggml-vulkan/vulkan-shaders/pool2d.comp | 74 + .../src/ggml-vulkan/vulkan-shaders/relu.comp | 21 + .../ggml-vulkan/vulkan-shaders/repeat.comp | 26 + .../ggml-vulkan/vulkan-shaders/rms_norm.comp | 42 + .../ggml-vulkan/vulkan-shaders/rope_head.comp | 49 + .../ggml-vulkan/vulkan-shaders/rope_neox.comp | 37 + .../ggml-vulkan/vulkan-shaders/rope_norm.comp | 37 + .../src/ggml-vulkan/vulkan-shaders/scale.comp | 24 + .../src/ggml-vulkan/vulkan-shaders/silu.comp | 22 + .../src/ggml-vulkan/vulkan-shaders/sin.comp | 17 + .../ggml-vulkan/vulkan-shaders/soft_max.comp | 174 + .../ggml-vulkan/vulkan-shaders/square.comp | 17 + .../ggml-vulkan/vulkan-shaders/sum_rows.comp | 37 + .../src/ggml-vulkan/vulkan-shaders/tanh.comp | 20 + .../vulkan-shaders/test_coopmat2_support.comp | 7 + .../vulkan-shaders/timestep_embedding.comp | 41 + .../src/ggml-vulkan/vulkan-shaders/types.comp | 323 + .../ggml-vulkan/vulkan-shaders/upscale.comp | 36 + .../vulkan-shaders/vulkan-shaders-gen.cpp | 594 ++ .../src/ggml-vulkan/vulkan-shaders/wkv6.comp | 87 + 76 files changed, 14642 insertions(+), 4 deletions(-) create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp diff --git a/CMakePresets.json b/CMakePresets.json index 442cb2a6d..09e924011 100644 --- a/CMakePresets.json +++ b/CMakePresets.json @@ -58,7 +58,11 @@ "cacheVariables": { "AMDGPU_TARGETS": "gfx900;gfx940;gfx941;gfx942;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102;gfx906:xnack-;gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-" } - } + }, + { + "name": "Vulkan", + "inherits": [ "Default" ] + } ], "buildPresets": [ { @@ -105,6 +109,11 @@ "name": "ROCm 6", "inherits": [ "ROCm" ], "configurePreset": "ROCm 6" - } + }, + { + "name": "Vulkan", + "targets": [ "ggml-vulkan" ], + "configurePreset": "Vulkan" + } ] } diff --git a/discover/gpu.go b/discover/gpu.go index 791d6b199..2494469a7 100644 --- a/discover/gpu.go +++ b/discover/gpu.go @@ -196,7 +196,10 @@ func initVulkanHandles() *vulkanHandles { libcapPaths := FindLibCapLibs() if len(vulkanPaths) > 0 && len(libcapPaths) > 0 { + slog.Info("vulkan: load libvulkan and libcap ok") vHandles.deviceCount, vHandles.vulkan, vulkanLibPath, libcapLibPath = LoadVulkanMgmt(vulkanPaths, libcapPaths) + } else { + slog.Info("vulkan: failed to load libvulkan or libcap") } return vHandles @@ -425,7 +428,7 @@ func GetGPUInfo() GpuInfoList { gpuInfo.ID = C.GoString(&memInfo.gpu_id[0]) gpuInfo.Compute = fmt.Sprintf("%d.%d", memInfo.major, memInfo.minor) gpuInfo.MinimumMemory = 0 - gpuInfo.DependencyPath = depPaths + gpuInfo.DependencyPath = []string{LibOllamaPath} gpuInfo.Name = C.GoString(&memInfo.gpu_name[0]) gpuInfo.DriverMajor = int(memInfo.major) gpuInfo.DriverMinor = int(memInfo.minor) @@ -767,7 +770,7 @@ func LoadVulkanMgmt(vulkanLibPaths []string, capLibPaths []string) (int, *C.vk_h C.vk_init(vkLib, capLib, &resp) if resp.err != nil { - slog.Debug("Unable to load vulkan", "library", vkLibPath, capLibPath, "error", C.GoString(resp.err)) + slog.Error("Unable to load vulkan", "library", vkLibPath, capLibPath, "error", C.GoString(resp.err)) C.free(unsafe.Pointer(resp.err)) } else { return int(resp.num_devices), &resp.ch, vkLibPath, capLibPath diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt b/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt new file mode 100644 index 000000000..9501de736 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt @@ -0,0 +1,92 @@ +find_package(Vulkan COMPONENTS glslc REQUIRED) + +if (Vulkan_FOUND) + message(STATUS "Vulkan found") + + ggml_add_backend_library(ggml-vulkan + ggml-vulkan.cpp + ../../include/ggml-vulkan.h + ) + + # Compile a test shader to determine whether GL_NV_cooperative_matrix2 is supported. + # If it's not, there will be an error to stderr. + # If it's supported, set a define to indicate that we should compile those shaders + execute_process(COMMAND ${Vulkan_GLSLC_EXECUTABLE} -o - -fshader-stage=compute --target-env=vulkan1.3 "${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders/test_coopmat2_support.comp" + OUTPUT_VARIABLE glslc_output + ERROR_VARIABLE glslc_error) + + if (${glslc_error} MATCHES ".*extension not supported: GL_NV_cooperative_matrix2.*") + message(STATUS "GL_NV_cooperative_matrix2 not supported by glslc") + else() + message(STATUS "GL_NV_cooperative_matrix2 supported by glslc") + add_compile_definitions(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + endif() + + target_link_libraries(ggml-vulkan PRIVATE Vulkan::Vulkan) + target_include_directories(ggml-vulkan PRIVATE ${CMAKE_CURRENT_BINARY_DIR}) + + # Workaround to the "can't dereference invalidated vector iterator" bug in clang-cl debug build + # Posssibly relevant: https://stackoverflow.com/questions/74748276/visual-studio-no-displays-the-correct-length-of-stdvector + if (MSVC AND CMAKE_CXX_COMPILER_ID STREQUAL "Clang") + add_compile_definitions(_ITERATOR_DEBUG_LEVEL=0) + endif() + + if (GGML_VULKAN_CHECK_RESULTS) + add_compile_definitions(GGML_VULKAN_CHECK_RESULTS) + endif() + + if (GGML_VULKAN_DEBUG) + add_compile_definitions(GGML_VULKAN_DEBUG) + endif() + + if (GGML_VULKAN_MEMORY_DEBUG) + add_compile_definitions(GGML_VULKAN_MEMORY_DEBUG) + endif() + + if (GGML_VULKAN_SHADER_DEBUG_INFO) + add_compile_definitions(GGML_VULKAN_SHADER_DEBUG_INFO) + endif() + + if (GGML_VULKAN_PERF) + add_compile_definitions(GGML_VULKAN_PERF) + endif() + + if (GGML_VULKAN_VALIDATE) + add_compile_definitions(GGML_VULKAN_VALIDATE) + endif() + + if (GGML_VULKAN_RUN_TESTS) + add_compile_definitions(GGML_VULKAN_RUN_TESTS) + endif() + + add_subdirectory(vulkan-shaders) + + set (_ggml_vk_genshaders_cmd vulkan-shaders-gen) + set (_ggml_vk_header ${CMAKE_CURRENT_BINARY_DIR}/ggml-vulkan-shaders.hpp) + set (_ggml_vk_source ${CMAKE_CURRENT_BINARY_DIR}/ggml-vulkan-shaders.cpp) + set (_ggml_vk_input_dir ${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders) + set (_ggml_vk_output_dir ${CMAKE_CURRENT_BINARY_DIR}/vulkan-shaders.spv) + + file(GLOB _ggml_vk_shader_deps "${_ggml_vk_input_dir}/*.comp") + + add_custom_command( + OUTPUT ${_ggml_vk_header} + ${_ggml_vk_source} + + COMMAND "$/${_ggml_vk_genshaders_cmd}" + --glslc ${Vulkan_GLSLC_EXECUTABLE} + --input-dir ${_ggml_vk_input_dir} + --output-dir ${_ggml_vk_output_dir} + --target-hpp ${_ggml_vk_header} + --target-cpp ${_ggml_vk_source} + --no-clean + + DEPENDS ${_ggml_vk_shader_deps} ${_ggml_vk_genshaders_cmd} + COMMENT "Generate vulkan shaders" + ) + + target_sources(ggml-vulkan PRIVATE ${_ggml_vk_source} ${_ggml_vk_header}) + +else() + message(WARNING "Vulkan not found") +endif() diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp new file mode 100644 index 000000000..d75cd6d61 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -0,0 +1,8745 @@ +#include "ggml-vulkan.h" +#include +#if defined(GGML_VULKAN_RUN_TESTS) || defined(GGML_VULKAN_PERF) || defined(GGML_VULKAN_CHECK_RESULTS) +#include +#include "ggml-cpu.h" +#endif + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ggml-impl.h" +#include "ggml-backend-impl.h" + +#include "ggml-vulkan-shaders.hpp" + +#define VK_API_VERSION VK_API_VERSION_1_2 + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +#define VK_VENDOR_ID_AMD 0x1002 +#define VK_VENDOR_ID_APPLE 0x106b +#define VK_VENDOR_ID_INTEL 0x8086 +#define VK_VENDOR_ID_NVIDIA 0x10de + +#define VK_DEVICE_DESCRIPTOR_POOL_SIZE 32 + +#define GGML_VK_MAX_NODES 8192 + +#define MAX_VK_BUFFERS 256 + +#define VK_CHECK(err, msg) \ + do { \ + vk::Result err_ = (err); \ + if (err_ != vk::Result::eSuccess) { \ + fprintf(stderr, "ggml_vulkan: %s error %s at %s:%d\n", \ + #err, to_string(err_).c_str(), __FILE__, __LINE__); \ + exit(1); \ + } \ + } while (0) + +#ifdef GGML_VULKAN_DEBUG +#define VK_LOG_DEBUG(msg) std::cerr << msg << std::endl +#else +#define VK_LOG_DEBUG(msg) ((void) 0) +#endif // GGML_VULKAN_DEBUG + +struct ggml_backend_vk_context; + +struct vk_queue { + uint32_t queue_family_index; + vk::Queue queue; + vk::CommandPool pool; + uint32_t cmd_buffer_idx; + std::vector cmd_buffers; + + vk::PipelineStageFlags stage_flags; + + bool transfer_only; +}; + +struct vk_pipeline_struct { + std::string name; + vk::ShaderModule shader_module; + vk::DescriptorSetLayout dsl; + std::vector descriptor_pools; + std::vector descriptor_sets; + uint32_t descriptor_set_idx; + vk::PipelineLayout layout; + vk::Pipeline pipeline; + uint32_t push_constant_size; + uint32_t parameter_count; + std::array wg_denoms; + uint32_t align; +}; + +typedef std::shared_ptr vk_pipeline; +typedef std::weak_ptr vk_pipeline_ref; + +static void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline); + +struct vk_matmul_pipeline_struct { + vk_pipeline l, m, s; + vk_pipeline a_l, a_m, a_s; +}; + +typedef std::shared_ptr vk_matmul_pipeline; + +struct vk_matmul_pipeline2 { + vk_matmul_pipeline2() { + f16acc = std::make_shared(); + f32acc = std::make_shared(); + } + vk_matmul_pipeline f32acc; + vk_matmul_pipeline f16acc; +}; + +struct vk_device_struct; +typedef std::shared_ptr vk_device; +typedef std::weak_ptr vk_device_ref; + +struct vk_buffer_struct; +typedef std::shared_ptr vk_buffer; +typedef std::weak_ptr vk_buffer_ref; + +struct ggml_backend_vk_buffer_type_context { + std::string name; + vk_device device; +}; + +static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft); +static ggml_backend_buffer_t ggml_backend_vk_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size); +static size_t ggml_backend_vk_buffer_type_get_alignment(ggml_backend_buffer_type_t buft); +static size_t ggml_backend_vk_buffer_type_get_max_size(ggml_backend_buffer_type_t buft); +static size_t ggml_backend_vk_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor); +static ggml_backend_buffer_type_i ggml_backend_vk_buffer_type_interface = { + /* .get_name = */ ggml_backend_vk_buffer_type_name, + /* .alloc_buffer = */ ggml_backend_vk_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_vk_buffer_type_get_alignment, + /* .get_max_size = */ ggml_backend_vk_buffer_type_get_max_size, + /* .get_alloc_size = */ ggml_backend_vk_buffer_type_get_alloc_size, + /* .is_host = */ NULL, +}; + +#ifdef GGML_VULKAN_MEMORY_DEBUG +class vk_memory_logger; +#endif +#ifdef GGML_VULKAN_PERF +class vk_perf_logger; +#endif +static void ggml_vk_destroy_buffer(vk_buffer& buf); + +static constexpr uint32_t mul_mat_vec_max_cols = 8; + +struct vk_device_struct { + std::mutex mutex; + + vk::PhysicalDevice physical_device; + vk::PhysicalDeviceProperties properties; + std::string name; + uint64_t max_memory_allocation_size; + bool fp16; + bool pipeline_robustness; + vk::Device device; + uint32_t vendor_id; + vk_queue compute_queue; + vk_queue transfer_queue; + bool single_queue; + uint32_t subgroup_size; + uint32_t shader_core_count; + bool uma; + bool float_controls_rte_fp16; + + bool subgroup_size_control; + uint32_t subgroup_min_size; + uint32_t subgroup_max_size; + bool subgroup_require_full_support; + + bool coopmat_support; + bool coopmat_acc_f32_support; + bool coopmat_acc_f16_support; + uint32_t coopmat_m; + uint32_t coopmat_n; + uint32_t coopmat_k; + bool coopmat2; + + size_t idx; + + bool mul_mat_l; + bool mul_mat_m; + bool mul_mat_s; + bool mul_mat_id_l; + bool mul_mat_id_m; + bool mul_mat_id_s; + + vk_matmul_pipeline pipeline_matmul_f32; + vk_matmul_pipeline pipeline_matmul_f32_f16; + vk_matmul_pipeline2 pipeline_matmul_f16; + vk_matmul_pipeline2 pipeline_matmul_f16_f32; + vk_pipeline pipeline_matmul_split_k_reduce; + + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_COUNT]; + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat[GGML_TYPE_COUNT]; + + vk_matmul_pipeline pipeline_matmul_id_f32; + vk_matmul_pipeline2 pipeline_matmul_id_f16; + vk_matmul_pipeline2 pipeline_matmul_id_f16_f32; + + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id[GGML_TYPE_COUNT]; + + vk_pipeline pipeline_dequant[GGML_TYPE_COUNT]; + vk_pipeline pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_COUNT][mul_mat_vec_max_cols]; + vk_pipeline pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_COUNT][mul_mat_vec_max_cols]; + vk_pipeline pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_COUNT]; + + vk_pipeline pipeline_mul_mat_vec_p021_f16_f32; + vk_pipeline pipeline_mul_mat_vec_nc_f16_f32; + vk_pipeline pipeline_get_rows[GGML_TYPE_COUNT]; + vk_pipeline pipeline_get_rows_f32[GGML_TYPE_COUNT]; + vk_pipeline pipeline_acc_f32; + vk_pipeline pipeline_add_f32, pipeline_add_f32_norepeat; + vk_pipeline pipeline_add_f16_f32_f16, pipeline_add_f16_f32_f16_norepeat; + vk_pipeline pipeline_mul_f32, pipeline_mul_f32_norepeat; + vk_pipeline pipeline_div_f32, pipeline_div_f32_norepeat; + vk_pipeline pipeline_concat_f32, pipeline_concat_f16, pipeline_concat_i32; + vk_pipeline pipeline_upscale_f32; + vk_pipeline pipeline_scale_f32; + vk_pipeline pipeline_sqr_f32; + vk_pipeline pipeline_sin_f32; + vk_pipeline pipeline_cos_f32; + vk_pipeline pipeline_clamp_f32; + vk_pipeline pipeline_pad_f32; + vk_pipeline pipeline_repeat_f32; + vk_pipeline pipeline_cpy_f32_f32, pipeline_cpy_f32_f16, pipeline_cpy_f16_f16; + vk_pipeline pipeline_contig_cpy_f32_f32, pipeline_contig_cpy_f32_f16, pipeline_contig_cpy_f16_f16; + vk_pipeline pipeline_norm_f32; + vk_pipeline pipeline_group_norm_f32; + vk_pipeline pipeline_rms_norm_f32; + vk_pipeline pipeline_gelu_f32; + vk_pipeline pipeline_gelu_quick_f32; + vk_pipeline pipeline_silu_f32; + vk_pipeline pipeline_relu_f32; + vk_pipeline pipeline_leaky_relu_f32; + vk_pipeline pipeline_tanh_f32; + vk_pipeline pipeline_diag_mask_inf_f32; + vk_pipeline pipeline_soft_max_f32, pipeline_soft_max_f32_f16; + vk_pipeline pipeline_soft_max_f32_wg512, pipeline_soft_max_f32_f16_wg512; + vk_pipeline pipeline_rope_norm_f32, pipeline_rope_norm_f16; + vk_pipeline pipeline_rope_neox_f32, pipeline_rope_neox_f16; + vk_pipeline pipeline_argsort_f32; + vk_pipeline pipeline_sum_rows_f32; + vk_pipeline pipeline_im2col_f32, pipeline_im2col_f32_f16; + vk_pipeline pipeline_timestep_embedding_f32; + vk_pipeline pipeline_pool2d_f32; + vk_pipeline pipeline_rwkv_wkv6_f32; + + // [2][2][2] is for {f16acc,f32acc}x{large,small_rows}x{unaligned, aligned} + vk_pipeline pipeline_flash_attn_f32_f16_D64[GGML_TYPE_COUNT][2][2][2]; + vk_pipeline pipeline_flash_attn_f32_f16_D80[GGML_TYPE_COUNT][2][2][2]; + vk_pipeline pipeline_flash_attn_f32_f16_D96[GGML_TYPE_COUNT][2][2][2]; + vk_pipeline pipeline_flash_attn_f32_f16_D112[GGML_TYPE_COUNT][2][2][2]; + vk_pipeline pipeline_flash_attn_f32_f16_D128[GGML_TYPE_COUNT][2][2][2]; + vk_pipeline pipeline_flash_attn_f32_f16_D256[GGML_TYPE_COUNT][2][2][2]; + + std::unordered_map pipelines; + std::unordered_map pipeline_descriptor_set_requirements; + + std::vector> pinned_memory; + + vk::Fence fence; + vk_buffer sync_staging; + + ggml_backend_buffer_type buffer_type; + +#ifdef GGML_VULKAN_MEMORY_DEBUG + std::unique_ptr memory_logger; +#endif +#ifdef GGML_VULKAN_PERF + std::unique_ptr perf_logger; +#endif + + ~vk_device_struct() { + VK_LOG_DEBUG("destroy device " << name); + + device.destroyFence(fence); + + ggml_vk_destroy_buffer(sync_staging); + + device.destroyCommandPool(compute_queue.pool); + if (!single_queue) { + device.destroyCommandPool(transfer_queue.pool); + } + + for (auto& pipeline : pipelines) { + if (pipeline.second.expired()) { + continue; + } + + vk_pipeline pl = pipeline.second.lock(); + ggml_vk_destroy_pipeline(device, pl); + } + pipelines.clear(); + + device.destroy(); + } +}; + +struct vk_buffer_struct { + vk::Buffer buffer = VK_NULL_HANDLE; + vk::DeviceMemory device_memory = VK_NULL_HANDLE; + vk::MemoryPropertyFlags memory_property_flags; + void * ptr; + size_t size = 0; + + vk_device device; + + ~vk_buffer_struct() { + if (size == 0) { + return; + } + VK_LOG_DEBUG("~vk_buffer_struct(" << buffer << ", " << size << ")"); + + device->device.freeMemory(device_memory); + device->device.destroyBuffer(buffer); + } +}; + +struct vk_subbuffer { + vk_buffer buffer; + uint64_t offset; + uint64_t size; + + operator vk::DescriptorBufferInfo() const { + return { buffer->buffer, offset, size }; + } +}; + +struct vk_semaphore { + vk::Semaphore s; + uint64_t value; +}; + +struct vk_submission { + vk::CommandBuffer buffer; + std::vector wait_semaphores; + std::vector signal_semaphores; +}; + +typedef std::vector vk_sequence; + +struct vk_mat_mat_push_constants { + uint32_t M; uint32_t N; uint32_t K; + uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; + uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; + uint32_t k_split; + uint32_t ne02; uint32_t ne12; uint32_t broadcast2; uint32_t broadcast3; +}; +struct vk_mat_vec_push_constants { + uint32_t ncols; uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; + uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; + uint32_t ne02; uint32_t ne12; uint32_t broadcast2; uint32_t broadcast3; +}; + +struct vk_mat_mat_id_push_constants { + uint32_t M; uint32_t N; uint32_t K; + uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; + uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; + uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11; +}; +struct vk_mat_vec_id_push_constants { + uint32_t ncols; uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; + uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; + uint32_t nei0; uint32_t ne11; +}; + +struct vk_flash_attn_push_constants { + uint32_t N; + uint32_t KV; + + uint32_t ne1; + uint32_t ne2; + uint32_t ne3; + + uint32_t neq2; + uint32_t neq3; + uint32_t nek2; + uint32_t nek3; + uint32_t nev2; + uint32_t nev3; + uint32_t nem1; + + uint32_t nb02; + uint32_t nb03; + uint32_t nb12; + uint32_t nb13; + uint32_t nb22; + uint32_t nb23; + uint32_t nb31; + + float scale; + float max_bias; + float logit_softcap; + + uint32_t mask; + uint32_t n_head_log2; + float m0; + float m1; +}; + +struct vk_op_push_constants { + uint32_t KX; + uint32_t KY; + float param1; + float param2; +}; + +struct vk_op_unary_push_constants { + uint32_t ne; + uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; + uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; uint32_t nb10; uint32_t nb11; uint32_t nb12; uint32_t nb13; + uint32_t misalign_offsets; + float param1; float param2; + uint32_t ne0_012mp; uint32_t ne0_012L; + uint32_t ne0_01mp; uint32_t ne0_01L; + uint32_t ne0_0mp; uint32_t ne0_0L; + uint32_t ne1_012mp; uint32_t ne1_012L; + uint32_t ne1_01mp; uint32_t ne1_01L; + uint32_t ne1_0mp; uint32_t ne1_0L; +}; +static_assert(sizeof(vk_op_unary_push_constants) <= 128, "sizeof(vk_op_unary_push_constants) must be <= 128"); + +// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1. +// Precompute mp (m' in the paper) and L such that division +// can be computed using a multiply (high 32b of 64b result) +// and a shift: +// +// n/d = (mulhi(n, mp) + n) >> L; +static void init_fastdiv_values(uint32_t d, uint32_t &mp, uint32_t &L) +{ + // compute L = ceil(log2(d)); + L = 0; + while (L < 32 && (uint32_t{1} << L) < d) { + L++; + } + + mp = (uint32_t)((uint64_t{1} << 32) * ((uint64_t{1} << L) - d) / d + 1); +} + +template void init_pushconst_fastdiv(T &p) { + GGML_UNUSED(p); + static_assert(!std::is_const::value, "unexpected type"); +} + +template <> void init_pushconst_fastdiv(vk_op_unary_push_constants &p) { + // Compute magic values to divide by these six numbers. + init_fastdiv_values(p.ne02*p.ne01*p.ne00, p.ne0_012mp, p.ne0_012L); + init_fastdiv_values(p.ne01*p.ne00, p.ne0_01mp, p.ne0_01L); + init_fastdiv_values(p.ne00, p.ne0_0mp, p.ne0_0L); + init_fastdiv_values(p.ne12*p.ne11*p.ne10, p.ne1_012mp, p.ne1_012L); + init_fastdiv_values(p.ne11*p.ne10, p.ne1_01mp, p.ne1_01L); + init_fastdiv_values(p.ne10, p.ne1_0mp, p.ne1_0L); +} + +struct vk_op_binary_push_constants { + uint32_t ne; + uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; + uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; uint32_t nb10; uint32_t nb11; uint32_t nb12; uint32_t nb13; + uint32_t ne20; uint32_t ne21; uint32_t ne22; uint32_t ne23; uint32_t nb20; uint32_t nb21; uint32_t nb22; uint32_t nb23; + uint32_t misalign_offsets; + float param1; float param2; int32_t param3; +}; + +struct vk_op_diag_mask_push_constants { + uint32_t ncols; + uint32_t rows_per_channel; + int32_t n_past; +}; + +struct vk_op_rope_push_constants { + uint32_t ncols; + uint32_t n_dims; + float freq_scale; + uint32_t p_delta_rows; + float freq_base; + float ext_factor; + float attn_factor; + float corr_dims[2]; + float theta_scale; + uint32_t has_ff; +}; + +struct vk_op_soft_max_push_constants { + uint32_t KX; + uint32_t KY; + float scale; + float max_bias; + float m0; + float m1; + uint32_t n_head_log2; + uint32_t nrows_x; +}; + +struct vk_op_argsort_push_constants { + uint32_t ncols; + uint32_t ncols_pad; + int32_t order; +}; + +struct vk_op_im2col_push_constants { + uint32_t batch_offset; uint32_t offset_delta; + uint32_t IC; + uint32_t IW; uint32_t IH; + uint32_t OW; uint32_t OH; + uint32_t KW; uint32_t KH; + uint32_t pelements; + uint32_t CHW; + int32_t s0; int32_t s1; + int32_t p0; int32_t p1; + int32_t d0; int32_t d1; +}; + +struct vk_op_timestep_embedding_push_constants { + uint32_t nb1; + uint32_t dim; + uint32_t max_period; +}; + +struct vk_op_pool2d_push_constants { + uint32_t IW; uint32_t IH; + uint32_t OW; uint32_t OH; + uint32_t OC; + uint32_t pelements; + uint32_t op; + int32_t k0; int32_t k1; + int32_t s0; int32_t s1; + int32_t p0; int32_t p1; +}; + +struct vk_op_rwkv_wkv6_push_constants { + uint32_t B; + uint32_t T; + uint32_t C; + uint32_t H; +}; + +// Allow pre-recording command buffers +struct vk_staging_memcpy { + vk_staging_memcpy(void * _dst, const void * _src, size_t _n) : dst(_dst), src(_src), n(_n) {} + + void * dst; + const void * src; + size_t n; +}; + +struct vk_op_upscale_push_constants { + uint32_t ne; uint32_t a_offset; uint32_t d_offset; + uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; + uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; + float sf0; float sf1; float sf2; float sf3; +}; + +struct vk_context_struct { + vk_submission * s; + std::vector seqs; + + int exit_tensor_idx; + + std::vector in_memcpys; + std::vector out_memcpys; + + vk_queue * q; +}; +typedef std::shared_ptr vk_context; +typedef std::weak_ptr vk_context_ref; + +struct ggml_vk_garbage_collector { + std::vector tl_semaphores; + std::vector semaphores; + std::vector events; + std::vector temp_buffers; + std::vector contexts; +}; + +#if defined(GGML_VULKAN_MEMORY_DEBUG) || defined(GGML_VULKAN_DEBUG) +#define VK_LOG_MEMORY(msg) std::cerr << "ggml_vulkan memory: " << msg << std::endl + +static std::string format_size(size_t size) { + const size_t kib = 1024; + const size_t mib = kib * 1024; + const size_t gib = mib * 1024; + + std::ostringstream oss; + oss << std::fixed << std::setprecision(2); + + if (size >= gib) { + oss << static_cast(size) / gib << " GiB"; + } else if (size >= mib) { + oss << static_cast(size) / mib << " MiB"; + } else if (size >= kib) { + oss << static_cast(size) / kib << " KiB"; + } else { + oss << size << " B"; + } + + return oss.str(); +} + +static std::mutex log_mutex; + +class vk_memory_logger { +public: + vk_memory_logger(): total_device(0), total_host(0) {} + void log_allocation(vk_buffer_ref buf_ref, size_t size); + void log_deallocation(vk_buffer_ref buf_ref); + +private: + std::map allocations; // Track allocations + size_t total_device; + size_t total_host; +}; +#else +#define VK_LOG_MEMORY(msg) ((void) 0) +#endif // GGML_VULKAN_MEMORY_DEBUG + +#if defined(GGML_VULKAN_PERF) + +class vk_perf_logger { +public: + void print_timings() { + std::cerr << "----------------\nVulkan Timings:" << std::endl; + for (const auto& t : timings) { + uint64_t total = 0; + for (const auto& time : t.second) { + total += time; + } + std::cerr << t.first << ": " << t.second.size() << " x " << (total / t.second.size() / 1000.0) << " ms" << std::endl; + } + + timings.clear(); + } + + void log_timing(const ggml_tensor * node, uint64_t time) { + if (node->op == GGML_OP_UNARY) { + timings[ggml_unary_op_name(ggml_get_unary_op(node))].push_back(time); + return; + } + if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { + const uint64_t m = node->src[0]->ne[1]; + const uint64_t n = node->src[1]->ne[1]; + const uint64_t k = node->src[1]->ne[0]; + std::string name = ggml_op_name(node->op); + if (n == 1) { + name += "_VEC m=" + std::to_string(m) + " k=" + std::to_string(k); + } else { + name += " m=" + std::to_string(m) + " n=" + std::to_string(n) + " k=" + std::to_string(k); + } + timings[name].push_back(time); + return; + } + timings[ggml_op_name(node->op)].push_back(time); + } +private: + std::map> timings; +}; +#endif // GGML_VULKAN_PERF + +struct ggml_backend_vk_context { + std::string name; + + vk_device device; + + size_t semaphore_idx, event_idx; + ggml_vk_garbage_collector gc; + size_t prealloc_size_x, prealloc_size_y, prealloc_size_split_k; + vk_buffer prealloc_x, prealloc_y, prealloc_split_k; + vk::Fence fence; + + vk_buffer buffer_pool[MAX_VK_BUFFERS]; + + vk_context_ref compute_ctx; + vk_context_ref transfer_ctx; + + std::vector tensor_ctxs; +}; + +static void * const vk_ptr_base = (void *)(uintptr_t) 0x1000; // NOLINT + +static uint64_t vk_tensor_offset(const ggml_tensor * tensor) { + if (tensor->view_src) { + return (uint8_t *) tensor->view_src->data - (uint8_t *) vk_ptr_base; + } + return (uint8_t *) tensor->data - (uint8_t *) vk_ptr_base; +} + +struct ggml_backend_vk_buffer_context { + vk_device_ref device; + vk_buffer dev_buffer; + std::string name; + + ggml_backend_vk_buffer_context(vk_device_ref device, vk_buffer&& dev_buffer, std::string& name) : + device(device), + dev_buffer(dev_buffer), + name(name) { + } + + ~ggml_backend_vk_buffer_context() { + ggml_vk_destroy_buffer(dev_buffer); + } +}; + +#ifdef GGML_VULKAN_MEMORY_DEBUG +void vk_memory_logger::log_allocation(vk_buffer_ref buf_ref, size_t size) { + std::lock_guard guard(log_mutex); + vk_buffer buf = buf_ref.lock(); + const bool device = bool(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eDeviceLocal); + const std::string type = device ? "device" : "host"; + allocations[buf->buffer] = size; + total_device += device ? size : 0; + total_host += device ? 0 : size; + VK_LOG_MEMORY(buf->device->name << ": +" << format_size(size) << " " << type << " at " << buf->buffer << ". Total device: " << format_size(total_device) << ", total host: " << format_size(total_host)); +} + +void vk_memory_logger::log_deallocation(vk_buffer_ref buf_ref) { + if (buf_ref.expired() || buf_ref.lock()->size == 0) { + return; + } + + std::lock_guard guard(log_mutex); + vk_buffer buf = buf_ref.lock(); + const bool device = bool(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eDeviceLocal); + std::string type = device ? "device" : "host"; + auto it = allocations.find(buf->buffer); + total_device -= device ? it->second : 0; + total_host -= device ? 0 : it->second; + if (it != allocations.end()) { + VK_LOG_MEMORY(buf->device->name << ": -" << format_size(it->second) << " " << type << " at " << buf->buffer << ". Total device: " << format_size(total_device) << ", total host: " << format_size(total_host)); + allocations.erase(it); + } else { + VK_LOG_MEMORY("ERROR " << buf->device->name << ": Attempted to deallocate unknown " << type << " memory at " << buf->buffer); + } +} +#endif // GGML_VULKAN_MEMORY_DEBUG + +struct vk_instance_t { + vk::Instance instance; + + std::vector device_indices; + vk_device devices[GGML_VK_MAX_DEVICES]; +}; + +static bool vk_instance_initialized = false; +static vk_instance_t vk_instance; + +#ifdef GGML_VULKAN_CHECK_RESULTS +static size_t vk_skip_checks; +static size_t vk_output_tensor; + +static void ggml_vk_print_tensor(const ggml_tensor * tensor, const char * name); +static void ggml_vk_check_results_0(ggml_tensor * tensor); +static void ggml_vk_check_results_1(ggml_tensor * tensor); +#endif + +typedef void (*ggml_vk_func_t)(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); + +static void ggml_backend_vk_free(ggml_backend_t backend); + +// variables to track number of compiles in progress +static uint32_t compile_count = 0; +static std::mutex compile_count_mutex; +static std::condition_variable compile_count_cond; + +static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipeline, const std::string name, size_t spv_size, const void* spv_data, const std::string entrypoint, + uint32_t parameter_count, uint32_t push_constant_size, std::array wg_denoms, std::vector specialization_constants, + uint32_t align, bool disable_robustness, bool require_full_subgroups, uint32_t required_subgroup_size) { + VK_LOG_DEBUG("ggml_vk_create_pipeline(" << device->name << ", " << name << ", " << entrypoint << ", " << parameter_count << ", " << push_constant_size << + ", (" << wg_denoms[0] << "," << wg_denoms[1] << "," << wg_denoms[2] << "), specialization_constants, " << align << + ", " << disable_robustness << ", " << require_full_subgroups << ", " << required_subgroup_size << ")"); + GGML_ASSERT(parameter_count > 0); + GGML_ASSERT(wg_denoms[0] > 0 && wg_denoms[1] > 0 && wg_denoms[2] > 0); // NOLINT + + pipeline = std::make_shared(); + pipeline->name = name; + pipeline->parameter_count = parameter_count; + pipeline->push_constant_size = push_constant_size; + pipeline->wg_denoms = wg_denoms; + pipeline->align = align; + + vk::ShaderModuleCreateInfo shader_module_create_info({}, spv_size, reinterpret_cast(spv_data)); + pipeline->shader_module = device->device.createShaderModule(shader_module_create_info); + + std::vector dsl_binding; + std::vector dsl_binding_flags; + for (uint32_t i = 0; i < parameter_count; i++) { + dsl_binding.push_back({i, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute}); + dsl_binding_flags.push_back({}); + } + + vk::DescriptorSetLayoutBindingFlagsCreateInfo dslbfci = { dsl_binding_flags }; + + vk::PushConstantRange pcr( + vk::ShaderStageFlagBits::eCompute, + 0, + pipeline->push_constant_size + ); + + vk::DescriptorSetLayoutCreateInfo descriptor_set_layout_create_info( + {}, + dsl_binding); + descriptor_set_layout_create_info.setPNext(&dslbfci); + pipeline->dsl = device->device.createDescriptorSetLayout(descriptor_set_layout_create_info); + + vk::DescriptorPoolSize descriptor_pool_size(vk::DescriptorType::eStorageBuffer, pipeline->parameter_count * VK_DEVICE_DESCRIPTOR_POOL_SIZE); + vk::DescriptorPoolCreateInfo descriptor_pool_create_info({}, VK_DEVICE_DESCRIPTOR_POOL_SIZE, descriptor_pool_size); + pipeline->descriptor_pools.push_back(device->device.createDescriptorPool(descriptor_pool_create_info)); + + pipeline->descriptor_set_idx = 0; + + vk::PipelineLayoutCreateInfo pipeline_layout_create_info(vk::PipelineLayoutCreateFlags(), pipeline->dsl, pcr); + pipeline->layout = device->device.createPipelineLayout(pipeline_layout_create_info); + + std::vector specialization_entries(specialization_constants.size()); + + for (size_t i = 0; i < specialization_constants.size(); i++) { + specialization_entries[i].constantID = i; + specialization_entries[i].offset = i * sizeof(uint32_t); + specialization_entries[i].size = sizeof(uint32_t); + } + + vk::SpecializationInfo specialization_info( + specialization_entries.size(), + specialization_entries.data(), + specialization_constants.size() * sizeof(uint32_t), + specialization_constants.data() + ); + + vk::PipelineShaderStageCreateFlags pipeline_shader_stage_create_flags{}; + + if (device->subgroup_require_full_support && require_full_subgroups) { + pipeline_shader_stage_create_flags |= vk::PipelineShaderStageCreateFlagBits::eRequireFullSubgroupsEXT; + } + + vk::PipelineShaderStageCreateInfo pipeline_shader_create_info( + pipeline_shader_stage_create_flags, + vk::ShaderStageFlagBits::eCompute, + pipeline->shader_module, + entrypoint.c_str(), + &specialization_info); + + vk::PipelineShaderStageRequiredSubgroupSizeCreateInfoEXT pipeline_shader_stage_required_subgroup_size_create_info; + pipeline_shader_stage_required_subgroup_size_create_info.requiredSubgroupSize = required_subgroup_size; + if (device->subgroup_size_control && required_subgroup_size > 0) { + GGML_ASSERT(device->subgroup_min_size <= required_subgroup_size && required_subgroup_size <= device->subgroup_max_size); + pipeline_shader_create_info.setPNext(&pipeline_shader_stage_required_subgroup_size_create_info); + } + + vk::ComputePipelineCreateInfo compute_pipeline_create_info( + vk::PipelineCreateFlags{}, + pipeline_shader_create_info, + pipeline->layout); + + vk::PipelineRobustnessCreateInfoEXT rci; + + if (device->pipeline_robustness && disable_robustness) { + rci.storageBuffers = vk::PipelineRobustnessBufferBehaviorEXT::eDisabled; + rci.uniformBuffers = vk::PipelineRobustnessBufferBehaviorEXT::eDisabled; + compute_pipeline_create_info.setPNext(&rci); + } + + pipeline->pipeline = device->device.createComputePipeline(VK_NULL_HANDLE, compute_pipeline_create_info).value; + + { + std::lock_guard guard(device->mutex); + device->pipelines.insert({ pipeline->name, pipeline }); + } + + { + std::lock_guard guard(compile_count_mutex); + assert(compile_count > 0); + compile_count--; + + // "Progress bar" for shader compiles + static uint32_t total_compile_count = 0; + if ((total_compile_count++ % 10) == 0) { + std::cerr << "."; + } + } + compile_count_cond.notify_all(); +} + +static void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline) { + VK_LOG_DEBUG("ggml_pipeline_destroy_pipeline(" << pipeline->name << ")"); + for (auto& pool : pipeline->descriptor_pools) { + device.destroyDescriptorPool(pool); + } + pipeline->descriptor_pools.clear(); + pipeline->descriptor_sets.clear(); + pipeline->descriptor_set_idx = 0; + + device.destroyDescriptorSetLayout(pipeline->dsl); + + device.destroyPipelineLayout(pipeline->layout); + + device.destroyShaderModule(pipeline->shader_module); + + device.destroyPipeline(pipeline->pipeline); +} + +static void ggml_pipeline_request_descriptor_sets(vk_device& device, vk_pipeline& pipeline, uint32_t n) { + VK_LOG_DEBUG("ggml_pipeline_request_descriptor_sets(" << pipeline->name << ", " << n << ")"); + device->pipeline_descriptor_set_requirements[pipeline->name] += n; +} + +static void ggml_pipeline_allocate_descriptor_sets(vk_device& device) { + std::lock_guard guard(device->mutex); + + for (auto& pair : device->pipeline_descriptor_set_requirements) { + vk_pipeline pipeline = device->pipelines.at(pair.first).lock(); + const uint64_t n = pair.second; + + VK_LOG_DEBUG("ggml_pipeline_allocate_descriptor_sets(" << pipeline->name << ", " << n << ")"); + + if (pipeline->descriptor_sets.size() >= pipeline->descriptor_set_idx + n) { + // Enough descriptors are available + continue; + } + + uint32_t to_alloc = pipeline->descriptor_set_idx + n - pipeline->descriptor_sets.size(); + uint32_t pool_remaining = VK_DEVICE_DESCRIPTOR_POOL_SIZE - pipeline->descriptor_sets.size() % VK_DEVICE_DESCRIPTOR_POOL_SIZE; + uint32_t pool_idx = pipeline->descriptor_sets.size() / VK_DEVICE_DESCRIPTOR_POOL_SIZE; + + while (to_alloc > 0) { + const uint32_t alloc_count = std::min(pool_remaining, to_alloc); + to_alloc -= alloc_count; + pool_remaining = VK_DEVICE_DESCRIPTOR_POOL_SIZE; + + if (pool_idx >= pipeline->descriptor_pools.size()) { + vk::DescriptorPoolSize descriptor_pool_size(vk::DescriptorType::eStorageBuffer, pipeline->parameter_count * VK_DEVICE_DESCRIPTOR_POOL_SIZE); + vk::DescriptorPoolCreateInfo descriptor_pool_create_info({}, VK_DEVICE_DESCRIPTOR_POOL_SIZE, descriptor_pool_size); + pipeline->descriptor_pools.push_back(device->device.createDescriptorPool(descriptor_pool_create_info)); + } + + std::vector layouts(alloc_count); + for (uint32_t i = 0; i < alloc_count; i++) { + layouts[i] = pipeline->dsl; + } + vk::DescriptorSetAllocateInfo descriptor_set_alloc_info(pipeline->descriptor_pools[pool_idx], alloc_count, layouts.data()); + std::vector sets = device->device.allocateDescriptorSets(descriptor_set_alloc_info); + pipeline->descriptor_sets.insert(pipeline->descriptor_sets.end(), sets.begin(), sets.end()); + + pool_idx++; + } + } +} + +static void ggml_pipeline_cleanup(vk_pipeline& pipeline) { + VK_LOG_DEBUG("ggml_pipeline_cleanup(" << pipeline->name << ")"); + pipeline->descriptor_set_idx = 0; +} + +static vk::CommandBuffer ggml_vk_create_cmd_buffer(vk_device& device, vk_queue& q) { + VK_LOG_DEBUG("ggml_vk_create_cmd_buffer()"); + std::lock_guard guard(device->mutex); + + if (q.cmd_buffers.size() > q.cmd_buffer_idx) { + // Reuse command buffer + return q.cmd_buffers[q.cmd_buffer_idx++]; + } + + vk::CommandBufferAllocateInfo command_buffer_alloc_info( + q.pool, + vk::CommandBufferLevel::ePrimary, + 1); + const std::vector cmd_buffers = device->device.allocateCommandBuffers(command_buffer_alloc_info); + auto buf = cmd_buffers.front(); + + q.cmd_buffers.push_back(buf); + q.cmd_buffer_idx++; + + return buf; +} + +static vk_submission ggml_vk_create_submission(vk_device& device, vk_queue& q, std::vector wait_semaphores, std::vector signal_semaphores) { + VK_LOG_DEBUG("ggml_vk_create_submission()"); + vk_submission s; + s.buffer = ggml_vk_create_cmd_buffer(device, q); + s.wait_semaphores = std::move(wait_semaphores); + s.signal_semaphores = std::move(signal_semaphores); + return s; +} + +static void ggml_vk_submit(vk_context& ctx, vk::Fence fence) { + if (ctx->seqs.empty()) { + if (fence) { + ctx->q->queue.submit({}, fence); + } + return; + } + VK_LOG_DEBUG("ggml_vk_submit(" << ctx << ", " << fence << ")"); + + std::vector> tl_wait_vals; + std::vector> tl_signal_vals; + std::vector> tl_wait_semaphores; + std::vector> tl_signal_semaphores; + std::vector tl_submit_infos; + std::vector submit_infos; + int idx = -1; + std::vector> stage_flags; + + size_t reserve = 0; + + for (const auto& sequence : ctx->seqs) { + reserve += sequence.size(); + } + + // Pre-reserve vectors to prevent reallocation, which invalidates pointers + tl_wait_semaphores.reserve(reserve); + tl_wait_vals.reserve(reserve); + tl_signal_semaphores.reserve(reserve); + tl_signal_vals.reserve(reserve); + tl_submit_infos.reserve(reserve); + submit_infos.reserve(reserve); + stage_flags.reserve(reserve); + + for (const auto& sequence : ctx->seqs) { + for (const auto& submission : sequence) { + stage_flags.push_back({}); + idx++; + tl_wait_vals.push_back({}); + tl_wait_semaphores.push_back({}); + tl_signal_vals.push_back({}); + tl_signal_semaphores.push_back({}); + for (size_t i = 0; i < submission.wait_semaphores.size(); i++) { + stage_flags[idx].push_back(ctx->q->stage_flags); + tl_wait_vals[idx].push_back(submission.wait_semaphores[i].value); + tl_wait_semaphores[idx].push_back(submission.wait_semaphores[i].s); + } + for (size_t i = 0; i < submission.signal_semaphores.size(); i++) { + tl_signal_vals[idx].push_back(submission.signal_semaphores[i].value); + tl_signal_semaphores[idx].push_back(submission.signal_semaphores[i].s); + } + tl_submit_infos.push_back({ + (uint32_t) submission.wait_semaphores.size(), + tl_wait_vals[idx].data(), + (uint32_t) submission.signal_semaphores.size(), + tl_signal_vals[idx].data(), + }); + tl_submit_infos[idx].sType = vk::StructureType::eTimelineSemaphoreSubmitInfo; + tl_submit_infos[idx].pNext = nullptr; + vk::SubmitInfo si{ + (uint32_t) submission.wait_semaphores.size(), + tl_wait_semaphores[idx].data(), + stage_flags[idx].data(), + 1, + &submission.buffer, + (uint32_t) submission.signal_semaphores.size(), + tl_signal_semaphores[idx].data(), + }; + si.setPNext(&tl_submit_infos[idx]); + submit_infos.push_back(si); + } + } + + ctx->q->queue.submit(submit_infos, fence); + + ctx->seqs.clear(); +} + +static uint32_t ggml_vk_find_queue_family_index(std::vector& queue_family_props, const vk::QueueFlags& required, const vk::QueueFlags& avoid, int32_t compute_index, uint32_t min_num_queues) { + VK_LOG_DEBUG("ggml_vk_find_queue_family_index()"); + const uint32_t qfsize = queue_family_props.size(); + + // Try with avoid preferences first + for (uint32_t i = 0; i < qfsize; i++) { + if (queue_family_props[i].queueCount >= min_num_queues && (compute_index < 0 || i != (uint32_t) compute_index) && queue_family_props[i].queueFlags & required && !(queue_family_props[i].queueFlags & avoid)) { + return i; + } + } + + // Fall back to only required + for (size_t i = 0; i < qfsize; i++) { + if (queue_family_props[i].queueCount >= min_num_queues && (compute_index < 0 || i != (uint32_t) compute_index) && queue_family_props[i].queueFlags & required) { + return i; + } + } + + // Fall back to reusing compute queue + for (size_t i = 0; i < qfsize; i++) { + if (queue_family_props[i].queueCount >= min_num_queues && queue_family_props[i].queueFlags & required) { + return i; + } + } + + // Fall back to ignoring min_num_queries + for (size_t i = 0; i < qfsize; i++) { + if (queue_family_props[i].queueFlags & required) { + return i; + } + } + + // All commands that are allowed on a queue that supports transfer operations are also allowed on a queue that supports either graphics or compute operations. + // Thus, if the capabilities of a queue family include VK_QUEUE_GRAPHICS_BIT or VK_QUEUE_COMPUTE_BIT, then reporting the VK_QUEUE_TRANSFER_BIT capability separately for that queue family is optional. + if (compute_index >= 0) { + return compute_index; + } + + std::cerr << "ggml_vulkan: No suitable queue family index found." << std::endl; + + for(auto &q_family : queue_family_props) { + std::cerr << "Queue number: " + std::to_string(q_family.queueCount) << " flags: " + to_string(q_family.queueFlags) << std::endl; + } + abort(); +} + +static void ggml_vk_create_queue(vk_device& device, vk_queue& q, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) { + VK_LOG_DEBUG("ggml_vk_create_queue()"); + std::lock_guard guard(device->mutex); + + q.queue_family_index = queue_family_index; + q.transfer_only = transfer_only; + + vk::CommandPoolCreateInfo command_pool_create_info_compute(vk::CommandPoolCreateFlags(VK_COMMAND_POOL_CREATE_TRANSIENT_BIT), queue_family_index); + q.pool = device->device.createCommandPool(command_pool_create_info_compute); + + q.cmd_buffer_idx = 0; + + q.queue = device->device.getQueue(queue_family_index, queue_index); + + q.stage_flags = stage_flags; +} + +static vk_context ggml_vk_create_context(ggml_backend_vk_context * ctx, vk_queue& q) { + vk_context result = std::make_shared(); + VK_LOG_DEBUG("ggml_vk_create_context(" << result << ")"); + ctx->gc.contexts.emplace_back(result); + result->q = &q; + return result; +} + +static vk_context ggml_vk_create_temporary_context(vk_queue& q) { + vk_context result = std::make_shared(); + VK_LOG_DEBUG("ggml_vk_create_temporary_context(" << result << ")"); + result->q = &q; + return result; +} + +static vk_semaphore * ggml_vk_create_binary_semaphore(ggml_backend_vk_context * ctx) { + VK_LOG_DEBUG("ggml_vk_create_timeline_semaphore()"); + vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eBinary, 0 }; + vk::SemaphoreCreateInfo ci{}; + ci.setPNext(&tci); + vk::Semaphore semaphore = ctx->device->device.createSemaphore(ci); + ctx->gc.semaphores.push_back({ semaphore, 0 }); + return &ctx->gc.semaphores[ctx->gc.semaphores.size() - 1]; +} + +static vk_semaphore * ggml_vk_create_timeline_semaphore(ggml_backend_vk_context * ctx) { + VK_LOG_DEBUG("ggml_vk_create_timeline_semaphore()"); + if (ctx->semaphore_idx >= ctx->gc.tl_semaphores.size()) { + vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eTimeline, 0 }; + vk::SemaphoreCreateInfo ci{}; + ci.setPNext(&tci); + vk::Semaphore semaphore = ctx->device->device.createSemaphore(ci); + ctx->gc.tl_semaphores.push_back({ semaphore, 0 }); + } + return &ctx->gc.tl_semaphores[ctx->semaphore_idx++]; +} + +static vk::Event ggml_vk_create_event(ggml_backend_vk_context * ctx) { + if (ctx->event_idx >= ctx->gc.events.size()) { + ctx->gc.events.push_back(ctx->device->device.createEvent({})); + } + return ctx->gc.events[ctx->event_idx++]; +} + +static void ggml_vk_queue_cleanup(vk_device& device, vk_queue& q) { + VK_LOG_DEBUG("ggml_vk_queue_cleanup()"); + std::lock_guard guard(device->mutex); + + // Requires command buffers to be done + device->device.resetCommandPool(q.pool); + q.cmd_buffer_idx = 0; +} + +static uint32_t find_properties(const vk::PhysicalDeviceMemoryProperties* mem_props, vk::MemoryRequirements* mem_req, vk::MemoryPropertyFlags flags) { + for (uint32_t i = 0; i < mem_props->memoryTypeCount; ++i) { + vk::MemoryType memory_type = mem_props->memoryTypes[i]; + if ((mem_req->memoryTypeBits & ((uint64_t)1 << i)) && + (flags & memory_type.propertyFlags) == flags && + mem_props->memoryHeaps[memory_type.heapIndex].size >= mem_req->size) { + return static_cast(i); + } + } + return UINT32_MAX; +} + +static vk_buffer ggml_vk_create_buffer(vk_device& device, size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags fallback_flags = vk::MemoryPropertyFlags(0)) { + VK_LOG_DEBUG("ggml_vk_create_buffer(" << device->name << ", " << size << ", " << to_string(req_flags) << ", " << to_string(fallback_flags) << ")"); + if (size > device->max_memory_allocation_size) { + throw vk::OutOfDeviceMemoryError("Requested buffer size exceeds device memory allocation limit"); + } + + std::lock_guard guard(device->mutex); + + vk_buffer buf = std::make_shared(); + + if (size == 0) { + buf->size = 0; + return buf; + } + + vk::BufferCreateInfo buffer_create_info{ + vk::BufferCreateFlags(), + size, + vk::BufferUsageFlagBits::eStorageBuffer | vk::BufferUsageFlagBits::eTransferSrc | vk::BufferUsageFlagBits::eTransferDst, + vk::SharingMode::eExclusive, + 0, + nullptr, + }; + + buf->buffer = device->device.createBuffer(buffer_create_info); + + vk::MemoryRequirements mem_req = device->device.getBufferMemoryRequirements(buf->buffer); + + vk::PhysicalDeviceMemoryProperties mem_props = device->physical_device.getMemoryProperties(); + + uint32_t memory_type_index = UINT32_MAX; + + memory_type_index = find_properties(&mem_props, &mem_req, req_flags); + buf->memory_property_flags = req_flags; + + if (memory_type_index == UINT32_MAX && fallback_flags) { + memory_type_index = find_properties(&mem_props, &mem_req, fallback_flags); + buf->memory_property_flags = fallback_flags; + } + + if (memory_type_index == UINT32_MAX) { + device->device.destroyBuffer(buf->buffer); + throw vk::OutOfDeviceMemoryError("No suitable memory type found"); + } + + try { + buf->device_memory = device->device.allocateMemory({ mem_req.size, memory_type_index }); + } catch (const vk::SystemError& e) { + if (buf->memory_property_flags != fallback_flags) { + // Try again with fallback flags + memory_type_index = find_properties(&mem_props, &mem_req, fallback_flags); + buf->memory_property_flags = fallback_flags; + + try { + buf->device_memory = device->device.allocateMemory({ mem_req.size, memory_type_index }); + } + catch (const vk::SystemError& e) { + device->device.destroyBuffer(buf->buffer); + throw e; + } + } else { + // Out of Host/Device memory, clean up buffer + device->device.destroyBuffer(buf->buffer); + throw e; + } + } + buf->ptr = nullptr; + + if (buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { + buf->ptr = device->device.mapMemory(buf->device_memory, 0, VK_WHOLE_SIZE); + } + + device->device.bindBufferMemory(buf->buffer, buf->device_memory, 0); + + buf->device = device; + buf->size = size; + +#ifdef GGML_VULKAN_MEMORY_DEBUG + device->memory_logger->log_allocation(buf, size); +#endif + + return buf; +} + +static vk_buffer ggml_vk_create_buffer_check(vk_device& device, size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags fallback_flags = vk::MemoryPropertyFlags(0)) { + try { + return ggml_vk_create_buffer(device, size, req_flags, fallback_flags); + } catch (const vk::SystemError& e) { + std::cerr << "ggml_vulkan: Memory allocation of size " << size << " failed." << std::endl; + std::cerr << "ggml_vulkan: " << e.what() << std::endl; + throw e; + } +} + +static vk_buffer ggml_vk_create_buffer_device(vk_device& device, size_t size) { + vk_buffer buf; + try { + if (device->uma) { + // Fall back to host memory type + buf = ggml_vk_create_buffer(device, size, vk::MemoryPropertyFlagBits::eDeviceLocal, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); + } else { + // use rebar if available, otherwise fallback to device only visible memory + buf = ggml_vk_create_buffer(device, size, vk::MemoryPropertyFlagBits::eDeviceLocal | vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent, vk::MemoryPropertyFlagBits::eDeviceLocal); + } + } catch (const vk::SystemError& e) { + std::cerr << "ggml_vulkan: Device memory allocation of size " << size << " failed." << std::endl; + std::cerr << "ggml_vulkan: " << e.what() << std::endl; + throw e; + } + + return buf; +} + +static void ggml_vk_destroy_buffer(vk_buffer& buf) { + if (buf == nullptr) { + return; + } + +#ifdef GGML_VULKAN_MEMORY_DEBUG + if (buf->device != nullptr) { + buf->device->memory_logger->log_deallocation(buf); + } +#endif + + buf.reset(); +} + +static vk_subbuffer ggml_vk_subbuffer(vk_buffer& buf) { + return { buf, 0, VK_WHOLE_SIZE }; +} + +static void ggml_vk_sync_buffers(vk_context& ctx) { + VK_LOG_DEBUG("ggml_vk_sync_buffers()"); + + const bool transfer_queue = ctx->q->transfer_only; + + ctx->s->buffer.pipelineBarrier( + ctx->q->stage_flags, + ctx->q->stage_flags, + {}, + { { + { !transfer_queue ? (vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) : (vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) }, + { !transfer_queue ? (vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) : (vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) } + } }, + {}, + {} + ); +} + +static void ggml_vk_wait_events(vk_context& ctx, std::vector&& events) { + VK_LOG_DEBUG("ggml_vk_wait_events()"); + if (events.empty()) { + return; + } + + ctx->s->buffer.waitEvents( + events, + ctx->q->stage_flags, + ctx->q->stage_flags, + {}, + {}, + {} + ); +} + +// number of rows/cols for flash attention shader +static constexpr uint32_t flash_attention_num_small_rows = 32; +static std::array fa_rows_cols(uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) { + GGML_UNUSED(clamp); + + // small rows, large cols + if (small_rows) { + return {flash_attention_num_small_rows, 128}; + } + // small cols to reduce register count + if (ggml_is_quantized(type) || D == 256) { + return {64, 32}; + } + return {64, 64}; +}; + +static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vector& warptile, bool mul_mat_id) { + // Needs to be kept up to date on shader changes + const uint32_t bank_conflict_offset = device->coopmat_support ? 8 : 1; + const uint32_t type_size = device->fp16 ? sizeof(ggml_fp16_t) : sizeof(float); + const uint32_t warps = warptile[0] / warptile[10]; + + const uint32_t load_bufs = (warptile[1] + warptile[2]) * (warptile[3] + bank_conflict_offset) * type_size; + const uint32_t mmid_row_ids = mul_mat_id ? 3072 * sizeof(uint32_t) : 0; + const uint32_t coopmat_stage = device->coopmat_support ? warptile[7] * warptile[8] / warps * sizeof(float) : 0; + + return (load_bufs + mmid_row_ids + coopmat_stage) <= device->properties.limits.maxComputeSharedMemorySize; +} + +static void ggml_vk_load_shaders(vk_device& device) { + VK_LOG_DEBUG("ggml_vk_load_shaders(" << device->name << ")"); + + std::cerr << "ggml_vulkan: Compiling shaders"; + + // some shaders have a minimum subgroup size + const uint32_t subgroup_size_16 = std::max(device->subgroup_size, 16u); + const uint32_t subgroup_size_32 = std::max(device->subgroup_size, 32u); + + // mulmat + std::vector l_warptile, m_warptile, s_warptile, + l_warptile_mmq, m_warptile_mmq, s_warptile_mmq, + l_warptile_mmq_k, m_warptile_mmq_k, s_warptile_mmq_k, + l_warptile_mmqid, m_warptile_mmqid, s_warptile_mmqid; + std::array l_wg_denoms, m_wg_denoms, s_wg_denoms, + l_mmq_wg_denoms, m_mmq_wg_denoms, s_mmq_wg_denoms, + l_mmq_wg_denoms_k, m_mmq_wg_denoms_k, s_mmq_wg_denoms_k, + l_mmqid_wg_denoms, m_mmqid_wg_denoms, s_mmqid_wg_denoms; + + uint32_t l_align, m_align, s_align; + if (device->coopmat2) { + // spec constants and tile sizes for non-quant matmul/matmul_id + l_warptile = { 256, 128, 256, 64 }; + m_warptile = { 256, 128, 128, 64 }; + s_warptile = { 128, 64, 64, 64 }; + l_wg_denoms = {128, 256, 1 }; + m_wg_denoms = {128, 128, 1 }; + s_wg_denoms = { 64, 64, 1 }; + + // spec constants and tile sizes for quant matmul (non-Qi_K) + l_warptile_mmq = { 256, 128, 256, 64 }; + m_warptile_mmq = { 256, 128, 128, 64 }; + s_warptile_mmq = { 256, 128, 128, 64 }; + l_mmq_wg_denoms = { 128, 256, 1 }; + m_mmq_wg_denoms = { 128, 128, 1 }; + s_mmq_wg_denoms = { 128, 128, 1 }; + + // spec constants and tile sizes for quant matmul (Qi_K) + l_warptile_mmq_k = { 256, 128, 512, 16 }; + m_warptile_mmq_k = { 256, 128, 256, 16 }; + s_warptile_mmq_k = { 256, 32, 128, 64 }; + l_mmq_wg_denoms_k = { 128, 512, 1 }; + m_mmq_wg_denoms_k = { 128, 256, 1 }; + s_mmq_wg_denoms_k = { 32, 128, 1 }; + + // spec constants and tile sizes for quant matmul_id + l_warptile_mmqid = { 256, 128, 128, 16 }; + m_warptile_mmqid = { 256, 128, 64, 16 }; + s_warptile_mmqid = { 256, 64, 64, 16 }; + l_mmqid_wg_denoms = { 128, 128, 1 }; + m_mmqid_wg_denoms = { 128, 64, 1 }; + s_mmqid_wg_denoms = { 64, 64, 1 }; + + l_align = 128; + m_align = 64; + s_align = 32; + } else { + // Matrix cores require different warp group sizes + const uint32_t tm_l = device->coopmat_support ? device->coopmat_m : 4; + const uint32_t tm_m = device->coopmat_support ? device->coopmat_m : 4; + const uint32_t tm_s = device->coopmat_support ? device->coopmat_m : 2; + const uint32_t tn_l = device->coopmat_support ? device->coopmat_n : 4; + const uint32_t tn_m = device->coopmat_support ? device->coopmat_n : 2; + const uint32_t tn_s = device->coopmat_support ? device->coopmat_n : 2; + const uint32_t tk_l = device->coopmat_support ? device->coopmat_k : 1; + const uint32_t tk_m = device->coopmat_support ? device->coopmat_k : 1; + const uint32_t tk_s = device->coopmat_support ? device->coopmat_k : 1; + + l_warptile = { 128, 128, 128, 16, device->subgroup_size * 2, 64, 2, tm_l, tn_l, tk_l, device->subgroup_size }; + m_warptile = { 128, 64, 64, 16, device->subgroup_size, 32, 2, tm_m, tn_m, tk_m, device->subgroup_size }; + s_warptile = { subgroup_size_16, 32, 32, 16, 32, 32, 2, tm_s, tn_s, tk_s, device->subgroup_size }; + + l_warptile_mmq = { 128, 128, 128, 32, device->subgroup_size * 2, 64, 2, tm_l, tn_l, tk_l, device->subgroup_size }; + m_warptile_mmq = { 128, 64, 64, 32, device->subgroup_size, 32, 2, tm_m, tn_m, tk_m, device->subgroup_size }; + s_warptile_mmq = { subgroup_size_32, 32, 32, 32, 32, 32, 2, tm_s, tn_s, tk_s, device->subgroup_size }; + + l_mmq_wg_denoms = l_wg_denoms = {128, 128, 1 }; + m_mmq_wg_denoms = m_wg_denoms = { 64, 64, 1 }; + s_mmq_wg_denoms = s_wg_denoms = { 32, 32, 1 }; + l_align = 128; + m_align = 64; + s_align = 32; + + // Fallback to smaller sizes if there's not enough shared memory. Given the current shaders + // and tile sizes, this should handle 16KB, 32KB, and 48KB+. + // This logic doesn't explicitly account for the 12KB row_ids in the mul_mat_mat_id shaders. + // But the numbers happen to work out for 32KB shared memory size that when using the medium + // size there's enough room for everything, and we assert for this. + uint32_t shmem_needed = (l_warptile[1] + l_warptile[2]) * (l_warptile[3] + 1) * sizeof(float); + if (shmem_needed > device->properties.limits.maxComputeSharedMemorySize) { + l_warptile = m_warptile; + l_wg_denoms = m_wg_denoms; + shmem_needed = (l_warptile[1] + l_warptile[2]) * (l_warptile[3] + 1) * sizeof(float); + GGML_ASSERT(shmem_needed <= device->properties.limits.maxComputeSharedMemorySize); + } + if (device->properties.limits.maxComputeSharedMemorySize >= 32768) { + // assert mul_mat_mat_id shaders will fit. + GGML_ASSERT(shmem_needed + 3072*4 <= device->properties.limits.maxComputeSharedMemorySize); + } + + shmem_needed = (l_warptile_mmq[1] + l_warptile_mmq[2]) * (l_warptile_mmq[3] + 1) * sizeof(float); + if (shmem_needed > device->properties.limits.maxComputeSharedMemorySize) { + if (device->properties.limits.maxComputeSharedMemorySize == 32768) { + l_warptile_mmq = m_warptile_mmq; + l_mmq_wg_denoms = m_mmq_wg_denoms; + } else { + l_warptile_mmq = s_warptile_mmq; + l_mmq_wg_denoms = s_mmq_wg_denoms; + } + shmem_needed = (l_warptile_mmq[1] + l_warptile_mmq[2]) * (l_warptile_mmq[3] + 1) * sizeof(float); + GGML_ASSERT(shmem_needed <= device->properties.limits.maxComputeSharedMemorySize); + } + if (device->properties.limits.maxComputeSharedMemorySize >= 32768) { + // assert mul_mat_mat_id shaders will fit. + GGML_ASSERT(shmem_needed + 3072*4 <= device->properties.limits.maxComputeSharedMemorySize); + } + // Disable medium and large matrix multiplication if not enough shared memory is available + // Check mmq warptiles as the largest configuration + // Throw an error if not enough for any matrix multiplication is available + if (!ggml_vk_matmul_shmem_support(device, s_warptile_mmq, false)) { + std::cerr << "ggml_vulkan: Error: Shared memory size too small for matrix multiplication." << std::endl; + throw std::runtime_error("Shared memory size too small for matrix multiplication."); + } else if (!ggml_vk_matmul_shmem_support(device, m_warptile_mmq, false)) { + device->mul_mat_m = false; + device->mul_mat_l = false; + } else if (!ggml_vk_matmul_shmem_support(device, l_warptile_mmq, false)) { + device->mul_mat_l = false; + } + + // Disable mul_mat_id if not enough shared memory is available + if (!ggml_vk_matmul_shmem_support(device, s_warptile_mmq, true)) { + device->mul_mat_id_s = false; + device->mul_mat_id_m = false; + device->mul_mat_id_l = false; + } else if (!ggml_vk_matmul_shmem_support(device, m_warptile_mmq, true)) { + device->mul_mat_id_m = false; + device->mul_mat_id_l = false; + } else if (!ggml_vk_matmul_shmem_support(device, l_warptile_mmq, true)) { + device->mul_mat_id_l = false; + } + } + + device->pipeline_matmul_f32 = std::make_shared(); + device->pipeline_matmul_f32_f16 = std::make_shared(); + + device->pipeline_matmul_id_f32 = std::make_shared(); + + std::vector> compiles; + auto const &ggml_vk_create_pipeline = [&](vk_device& device, vk_pipeline& pipeline, const std::string &name, size_t spv_size, const void* spv_data, const std::string &entrypoint, + uint32_t parameter_count, uint32_t push_constant_size, std::array wg_denoms, const std::vector& specialization_constants, + uint32_t align, bool disable_robustness = false, bool require_full_subgroups = false, uint32_t required_subgroup_size = 0) { + { + // wait until fewer than N compiles are in progress + uint32_t N = std::max(1u, std::thread::hardware_concurrency()); + std::unique_lock guard(compile_count_mutex); + while (compile_count >= N) { + compile_count_cond.wait(guard); + } + compile_count++; + } + compiles.push_back(std::async(ggml_vk_create_pipeline_func, std::ref(device), std::ref(pipeline), name, spv_size, spv_data, entrypoint, + parameter_count, push_constant_size, wg_denoms, specialization_constants, align, disable_robustness, require_full_subgroups, required_subgroup_size)); + }; + +#if defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + if (device->coopmat2) { + + auto const &fa_wg_denoms = [&](uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) -> std::array { + return {fa_rows_cols(D, clamp, type, small_rows)[0], 1, 1}; + }; + + auto const &fa_spec_constants = [&](uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) -> std::vector { + // For large number of rows, 128 invocations seems to work best. + // For small number of rows (e.g. N==1), 256 works better. But matrix granularity for 256 is 32, so we + // can't use 256 for D==80. + uint32_t wg_size = (small_rows && (D % 32) == 0) ? 256 : 128; + auto rows_cols = fa_rows_cols(D, clamp, type, small_rows); + return {wg_size, rows_cols[0], rows_cols[1], (D), clamp}; + }; + +#define CREATE_FA2(TYPE, NAMELC, D) \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][0][0], "flash_attn_f32_f16_D" #D "_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,false), fa_spec_constants(D,1,TYPE,false), 1); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][0][1], "flash_attn_f32_f16_D" #D "_aligned_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,false), fa_spec_constants(D,0,TYPE,false), fa_rows_cols(D,0,TYPE,false)[1]); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][0][0], "flash_attn_f32_f16_D" #D "_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,false), fa_spec_constants(D,1,TYPE,false), 1); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][0][1], "flash_attn_f32_f16_D" #D "_aligned_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,false), fa_spec_constants(D,0,TYPE,false), fa_rows_cols(D,0,TYPE,false)[1]); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][1][0], "flash_attn_f32_f16_D" #D "_f16acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,true), fa_spec_constants(D,1,TYPE,true), 1); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][1][1], "flash_attn_f32_f16_D" #D "_aligned_f16acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,true), fa_spec_constants(D,0,TYPE,true), fa_rows_cols(D,0,TYPE,true)[1]); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][1][0], "flash_attn_f32_f16_D" #D "_f32acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,true), fa_spec_constants(D,1,TYPE,true), 1); \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][1][1], "flash_attn_f32_f16_D" #D "_aligned_f32acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,true), fa_spec_constants(D,0,TYPE,true), fa_rows_cols(D,0,TYPE,true)[1]); \ + +#define CREATE_FA(TYPE, NAMELC) \ + CREATE_FA2(TYPE, NAMELC, 64) \ + CREATE_FA2(TYPE, NAMELC, 80) \ + CREATE_FA2(TYPE, NAMELC, 96) \ + CREATE_FA2(TYPE, NAMELC, 112) \ + CREATE_FA2(TYPE, NAMELC, 128) \ + CREATE_FA2(TYPE, NAMELC, 256) + + CREATE_FA(GGML_TYPE_F16, f16) + CREATE_FA(GGML_TYPE_Q4_0, q4_0) + CREATE_FA(GGML_TYPE_Q4_1, q4_1) + CREATE_FA(GGML_TYPE_Q5_0, q5_0) + CREATE_FA(GGML_TYPE_Q5_1, q5_1) + CREATE_FA(GGML_TYPE_Q8_0, q8_0) + // K dequants currently disabled because D dimension is rounded up to 256 and runs inefficiently + //CREATE_FA(GGML_TYPE_Q2_K, q2_k) + //CREATE_FA(GGML_TYPE_Q3_K, q3_k) + //CREATE_FA(GGML_TYPE_Q4_K, q4_k) + //CREATE_FA(GGML_TYPE_Q5_K, q5_k) + //CREATE_FA(GGML_TYPE_Q6_K, q6_k) + CREATE_FA(GGML_TYPE_IQ4_NL, iq4_nl) +#undef CREATE_FA + + // Create 6 variants, {s,m,l}x{unaligned,aligned} +#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ + + // Create 2 variants, {f16,f32} accumulator +#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + + CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3) + + CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + + CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) + CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) + CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) +#undef CREATE_MM +#undef CREATE_MM2 + } else +#endif // defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + if (device->coopmat_support) { + // Create 6 variants, {s,m,l}x{unaligned,aligned} +#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1, false, true); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1, false, true); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1, false, true); \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align, false, true); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align, false, true); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align, false, true); \ + + // Create 2 variants, {f16,f32} accumulator +#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->coopmat_acc_f16_support) { \ + CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + } \ + if (device->coopmat_acc_f32_support) { \ + CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + } \ + + CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + + if (device->coopmat_acc_f16_support) { + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + } else { + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + } + + // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. + if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { + CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + + if (device->coopmat_acc_f16_support) { + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + } else { + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + } + } +#undef CREATE_MM2 +#undef CREATE_MM + } else if (device->fp16) { + // Create 6 variants, {s,m,l}x{unaligned,aligned} +#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ + + // Create 2 variants, {f16,f32} accumulator +#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + + CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. + if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { + CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + } +#undef CREATE_MM2 +#undef CREATE_MM + } else { + // Create 6 variants, {s,m,l}x{unaligned,aligned} +#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _l) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ + if (device->mul_mat ## ID ## _m) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ + if (device->mul_mat ## ID ## _s) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ + + CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_matmul_f16.f32acc, matmul_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_matmul_f16_f32.f32acc, matmul_f16_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f32acc, matmul_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f32acc, matmul_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f32acc, matmul_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f32acc, matmul_q5_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f32acc, matmul_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f32acc, matmul_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f32acc, matmul_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f32acc, matmul_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f32acc, matmul_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f32acc, matmul_q6_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f32acc, matmul_iq4_nl_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. + if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { + CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM(pipeline_matmul_id_f16.f32acc, matmul_id_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + CREATE_MM(pipeline_matmul_id_f16_f32.f32acc, matmul_id_f16_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f32acc, matmul_id_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f32acc, matmul_id_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f32acc, matmul_id_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f32acc, matmul_id_q5_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f32acc, matmul_id_q6_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f32acc, matmul_id_iq4_nl_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); + } +#undef CREATE_MM + } + + // mul mat vec + + // the number of rows computed per shader depends on GPU model and quant + uint32_t rm_stdq = 1; + uint32_t rm_kq = 2; + if (device->vendor_id == VK_VENDOR_ID_AMD) { + if (device->subgroup_min_size == 64 && device->subgroup_max_size == 64) { // GCN + rm_stdq = 2; + rm_kq = 4; + } + } else if (device->vendor_id == VK_VENDOR_ID_INTEL) + rm_stdq = 2; + + for (uint32_t i = 0; i < mul_mat_vec_max_cols; ++i) { + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_F32 ][i], "mul_mat_vec_f32_f32_f32_"+std::to_string(i+1), mul_mat_vec_f32_f32_f32_len, mul_mat_vec_f32_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f32_f32_"+std::to_string(i+1), mul_mat_vec_f16_f32_f32_len, mul_mat_vec_f16_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_0_f32_f32_len, mul_mat_vec_q4_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_1_f32_f32_len, mul_mat_vec_q4_1_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_0_f32_f32_len, mul_mat_vec_q5_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_1_f32_f32_len, mul_mat_vec_q5_1_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q8_0_f32_f32_len, mul_mat_vec_q8_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q2_k_f32_f32_len, mul_mat_vec_q2_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q3_k_f32_f32_len, mul_mat_vec_q3_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_k_f32_f32_len, mul_mat_vec_q4_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_k_f32_f32_len, mul_mat_vec_q5_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q6_k_f32_f32_len, mul_mat_vec_q6_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f32_f32_"+std::to_string(i+1), mul_mat_vec_iq4_nl_f32_f32_len, mul_mat_vec_iq4_nl_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq, i+1}, 1, true); + + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_F32 ][i], "mul_mat_vec_f32_f16_f32_"+std::to_string(i+1), mul_mat_vec_f32_f16_f32_len, mul_mat_vec_f32_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f16_f32_"+std::to_string(i+1), mul_mat_vec_f16_f16_f32_len, mul_mat_vec_f16_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_0_f16_f32_len, mul_mat_vec_q4_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_1_f16_f32_len, mul_mat_vec_q4_1_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_0_f16_f32_len, mul_mat_vec_q5_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_1_f16_f32_len, mul_mat_vec_q5_1_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q8_0_f16_f32_len, mul_mat_vec_q8_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q2_k_f16_f32_len, mul_mat_vec_q2_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q3_k_f16_f32_len, mul_mat_vec_q3_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_k_f16_f32_len, mul_mat_vec_q4_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_k_f16_f32_len, mul_mat_vec_q5_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q6_k_f16_f32_len, mul_mat_vec_q6_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f16_f32_"+std::to_string(i+1), mul_mat_vec_iq4_nl_f16_f32_len, mul_mat_vec_iq4_nl_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq, i+1}, 1, true); + } + + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_F32 ], "mul_mat_vec_id_f32_f32", mul_mat_vec_id_f32_f32_len, mul_mat_vec_id_f32_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {device->subgroup_size, 2}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_F16 ], "mul_mat_vec_id_f16_f32", mul_mat_vec_id_f16_f32_len, mul_mat_vec_id_f16_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {device->subgroup_size, 2}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_0], "mul_mat_vec_id_q4_0_f32", mul_mat_vec_id_q4_0_f32_len, mul_mat_vec_id_q4_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_1], "mul_mat_vec_id_q4_1_f32", mul_mat_vec_id_q4_1_f32_len, mul_mat_vec_id_q4_1_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_0], "mul_mat_vec_id_q5_0_f32", mul_mat_vec_id_q5_0_f32_len, mul_mat_vec_id_q5_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_1], "mul_mat_vec_id_q5_1_f32", mul_mat_vec_id_q5_1_f32_len, mul_mat_vec_id_q5_1_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q8_0], "mul_mat_vec_id_q8_0_f32", mul_mat_vec_id_q8_0_f32_len, mul_mat_vec_id_q8_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q2_K], "mul_mat_vec_id_q2_k_f32", mul_mat_vec_id_q2_k_f32_len, mul_mat_vec_id_q2_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q3_K], "mul_mat_vec_id_q3_k_f32", mul_mat_vec_id_q3_k_f32_len, mul_mat_vec_id_q3_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_K], "mul_mat_vec_id_q4_k_f32", mul_mat_vec_id_q4_k_f32_len, mul_mat_vec_id_q4_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_K], "mul_mat_vec_id_q5_k_f32", mul_mat_vec_id_q5_k_f32_len, mul_mat_vec_id_q5_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q6_K], "mul_mat_vec_id_q6_k_f32", mul_mat_vec_id_q6_k_f32_len, mul_mat_vec_id_q6_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_IQ4_NL], "mul_mat_vec_id_iq4_nl_f32", mul_mat_vec_id_iq4_nl_f32_len, mul_mat_vec_id_iq4_nl_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq}, 1, true); + + // dequant shaders + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_F32 ], "f32_to_f16", dequant_f32_len, dequant_f32_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_0], "dequant_q4_0", dequant_q4_0_len, dequant_q4_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_1], "dequant_q4_1", dequant_q4_1_len, dequant_q4_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_1], "dequant_q5_1", dequant_q5_1_len, dequant_q5_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_K], "dequant_q4_k", dequant_q4_k_len, dequant_q4_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_K], "dequant_q5_k", dequant_q5_k_len, dequant_q5_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q6_K], "dequant_q6_k", dequant_q6_k_len, dequant_q6_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ4_NL], "dequant_iq4_nl", dequant_iq4_nl_len, dequant_iq4_nl_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + + // get_rows + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_F32 ], "get_rows_f32", get_rows_f32_len, get_rows_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_F16 ], "get_rows_f16", get_rows_f16_len, get_rows_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_0], "get_rows_q4_0", get_rows_q4_0_len, get_rows_q4_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_1], "get_rows_q4_1", get_rows_q4_1_len, get_rows_q4_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_0], "get_rows_q5_0", get_rows_q5_0_len, get_rows_q5_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_1], "get_rows_q5_1", get_rows_q5_1_len, get_rows_q5_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q8_0], "get_rows_q8_0", get_rows_q8_0_len, get_rows_q8_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_IQ4_NL], "get_rows_iq4_nl", get_rows_iq4_nl_len, get_rows_iq4_nl_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_F32 ], "get_rows_f32_f32", get_rows_f32_f32_len, get_rows_f32_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_F16 ], "get_rows_f16_f32", get_rows_f16_f32_len, get_rows_f16_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_0], "get_rows_q4_0_f32", get_rows_q4_0_f32_len, get_rows_q4_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_1], "get_rows_q4_1_f32", get_rows_q4_1_f32_len, get_rows_q4_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_0], "get_rows_q5_0_f32", get_rows_q5_0_f32_len, get_rows_q5_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_1], "get_rows_q5_1_f32", get_rows_q5_1_f32_len, get_rows_q5_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q8_0], "get_rows_q8_0_f32", get_rows_q8_0_f32_len, get_rows_q8_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_IQ4_NL], "get_rows_iq4_nl_f32", get_rows_iq4_nl_f32_len, get_rows_iq4_nl_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_matmul_split_k_reduce, "split_k_reduce", split_k_reduce_len, split_k_reduce_data, "main", 2, 2 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_mul_mat_vec_p021_f16_f32, "mul_mat_vec_p021_f16_f32", mul_mat_vec_p021_f16_f32_len, mul_mat_vec_p021_f16_f32_data, "main", 3, 6 * sizeof(uint32_t), {1, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_mul_mat_vec_nc_f16_f32, "mul_mat_vec_nc_f16_f32", mul_mat_vec_nc_f16_f32_len, mul_mat_vec_nc_f16_f32_data, "main", 3, 7 * sizeof(uint32_t), {1, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_norm_f32, "norm_f32", norm_f32_len, norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_group_norm_f32, "group_norm_f32", group_norm_f32_len, group_norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_f32, "rms_norm_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f32, "cpy_f32_f32", cpy_f32_f32_len, cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f16, "cpy_f32_f16", cpy_f32_f16_len, cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_f16_f16, "cpy_f16_f16", cpy_f16_f16_len, cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f32, "contig_cpy_f32_f32", contig_cpy_f32_f32_len, contig_cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f16, "contig_cpy_f32_f16", contig_cpy_f32_f16_len, contig_cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f16_f16, "contig_cpy_f16_f16", contig_cpy_f16_f16_len, contig_cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_add_f32, "add_f32", add_f32_len, add_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); + ggml_vk_create_pipeline(device, device->pipeline_add_f32_norepeat, "add_f32_norepeat", add_f32_len, add_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_add_f16_f32_f16, "add_f16_f32_f16", add_f16_f32_f16_len, add_f16_f32_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); + ggml_vk_create_pipeline(device, device->pipeline_add_f16_f32_f16_norepeat, "add_f16_f32_f16_norepeat", add_f16_f32_f16_len, add_f16_f32_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_acc_f32, "acc_f32", acc_f32_len, acc_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_mul_f32, "mul_f32", mul_f32_len, mul_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); + ggml_vk_create_pipeline(device, device->pipeline_mul_f32_norepeat, "mul_f32_norepeat", mul_f32_len, mul_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); + ggml_vk_create_pipeline(device, device->pipeline_div_f32, "div_f32", div_f32_len, div_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); + ggml_vk_create_pipeline(device, device->pipeline_div_f32_norepeat, "div_f32_norepeat", div_f32_len, div_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_concat_f32, "concat_f32", concat_f32_len, concat_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_concat_f16, "concat_f16", concat_f16_len, concat_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_concat_i32, "concat_i32", concat_i32_len, concat_i32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_upscale_f32, "upscale_f32", upscale_f32_len, upscale_f32_data, "main", 2, sizeof(vk_op_upscale_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_scale_f32, "scale_f32", scale_f32_len, scale_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_sqr_f32, "sqr_f32", sqr_f32_len, sqr_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_sin_f32, "sin_f32", sin_f32_len, sin_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cos_f32, "cos_f32", cos_f32_len, cos_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_clamp_f32, "clamp_f32", clamp_f32_len, clamp_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_pad_f32, "pad_f32", pad_f32_len, pad_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_repeat_f32, "repeat_f32", repeat_f32_len, repeat_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_gelu_f32, "gelu_f32", gelu_f32_len, gelu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_gelu_quick_f32, "gelu_quick_f32", gelu_quick_f32_len, gelu_quick_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_silu_f32, "silu_f32", silu_f32_len, silu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_relu_f32, "relu_f32", relu_f32_len, relu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_leaky_relu_f32, "leaky_relu_f32", leaky_relu_f32_len, leaky_relu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_tanh_f32, "tanh_f32", tanh_f32_len, tanh_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_diag_mask_inf_f32, "diag_mask_inf_f32", diag_mask_inf_f32_len, diag_mask_inf_f32_data, "main", 2, sizeof(vk_op_diag_mask_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32, "soft_max_f32", soft_max_f32_len, soft_max_f32_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); + ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_wg512, "soft_max_f32_wg512", soft_max_f32_len, soft_max_f32_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 512 }, 1); + ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_f16, "soft_max_f32_f16", soft_max_f32_f16_len, soft_max_f32_f16_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); + ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_f16_wg512, "soft_max_f32_f16_wg512", soft_max_f32_f16_len, soft_max_f32_f16_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 512 }, 1); + + ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f32, "rope_norm_f32", rope_norm_f32_len, rope_norm_f32_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f32, "rope_neox_f32", rope_neox_f32_len, rope_neox_f32_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + + if (device->float_controls_rte_fp16) { + ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f16, "rope_norm_f16", rope_norm_f16_rte_len, rope_norm_f16_rte_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f16, "rope_neox_f16", rope_neox_f16_rte_len, rope_neox_f16_rte_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + } else { + ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f16, "rope_norm_f16", rope_norm_f16_len, rope_norm_f16_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f16, "rope_neox_f16", rope_neox_f16_len, rope_neox_f16_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); + } + + ggml_vk_create_pipeline(device, device->pipeline_argsort_f32, "argsort_f32", argsort_f32_len, argsort_f32_data, "main", 2, sizeof(vk_op_argsort_push_constants), {1024, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); + + ggml_vk_create_pipeline(device, device->pipeline_im2col_f32, "im2col_f32", im2col_f32_len, im2col_f32_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); + if (device->float_controls_rte_fp16) { + ggml_vk_create_pipeline(device, device->pipeline_im2col_f32_f16, "im2col_f32_f16", im2col_f32_f16_rte_len, im2col_f32_f16_rte_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); + } else { + ggml_vk_create_pipeline(device, device->pipeline_im2col_f32_f16, "im2col_f32_f16", im2col_f32_f16_len, im2col_f32_f16_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); + } + + ggml_vk_create_pipeline(device, device->pipeline_timestep_embedding_f32, "timestep_embedding_f32", timestep_embedding_f32_len, timestep_embedding_f32_data, "main", 2, sizeof(vk_op_timestep_embedding_push_constants), {256, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_pool2d_f32, "pool2d_f32", pool2d_f32_len, pool2d_f32_data, "main", 2, sizeof(vk_op_pool2d_push_constants), {512, 1, 1}, {}, 1); + + ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv6_f32, "rwkv_wkv6_f32", rwkv_wkv6_f32_len, rwkv_wkv6_f32_data, "main", 7, sizeof(vk_op_rwkv_wkv6_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); + + for (auto &c : compiles) { + c.wait(); + } + std::cerr << "Done!" << std::endl; +} + +static bool ggml_vk_khr_cooperative_matrix_support(const vk::PhysicalDeviceProperties& props, const vk::PhysicalDeviceDriverProperties& driver_props); + +static vk_device ggml_vk_get_device(size_t idx) { + VK_LOG_DEBUG("ggml_vk_get_device(" << idx << ")"); + + if (vk_instance.devices[idx] == nullptr) { + VK_LOG_DEBUG("Initializing new vk_device"); + vk_device device = std::make_shared(); + vk_instance.devices[idx] = device; + +#ifdef GGML_VULKAN_MEMORY_DEBUG + device->memory_logger = std::unique_ptr(new vk_memory_logger()); +#endif +#ifdef GGML_VULKAN_PERF + device->perf_logger = std::unique_ptr(new vk_perf_logger()); +#endif + + size_t dev_num = vk_instance.device_indices[idx]; + + std::vector physical_devices = vk_instance.instance.enumeratePhysicalDevices(); + + if (dev_num >= physical_devices.size()) { + std::cerr << "ggml_vulkan: Device with index " << dev_num << " does not exist." << std::endl; + throw std::runtime_error("Device not found"); + } + + device->physical_device = physical_devices[dev_num]; + const std::vector ext_props = device->physical_device.enumerateDeviceExtensionProperties(); + + bool fp16_storage = false; + bool fp16_compute = false; + bool maintenance4_support = false; + bool sm_builtins = false; + bool amd_shader_core_properties2 = false; + bool pipeline_robustness = false; + bool coopmat2_support = false; + device->coopmat_support = false; + + // Check if maintenance4 is supported + for (const auto& properties : ext_props) { + if (strcmp("VK_KHR_maintenance4", properties.extensionName) == 0) { + maintenance4_support = true; + } else if (strcmp("VK_KHR_16bit_storage", properties.extensionName) == 0) { + fp16_storage = true; + } else if (strcmp("VK_KHR_shader_float16_int8", properties.extensionName) == 0) { + fp16_compute = true; + } else if (strcmp("VK_NV_shader_sm_builtins", properties.extensionName) == 0) { + sm_builtins = true; + } else if (strcmp("VK_AMD_shader_core_properties2", properties.extensionName) == 0) { + amd_shader_core_properties2 = true; + } else if (strcmp("VK_EXT_pipeline_robustness", properties.extensionName) == 0) { + pipeline_robustness = true; + } else if (strcmp("VK_EXT_subgroup_size_control", properties.extensionName) == 0) { + device->subgroup_size_control = true; + } else if (strcmp("VK_KHR_cooperative_matrix", properties.extensionName) == 0 && + !getenv("GGML_VK_DISABLE_COOPMAT")) { + device->coopmat_support = true; + device->coopmat_m = 0; + device->coopmat_n = 0; + device->coopmat_k = 0; + } else if (strcmp("VK_NV_cooperative_matrix2", properties.extensionName) == 0 && + !getenv("GGML_VK_DISABLE_COOPMAT2")) { + coopmat2_support = true; + } + } + + vk::PhysicalDeviceProperties2 props2; + vk::PhysicalDeviceMaintenance3Properties props3; + vk::PhysicalDeviceMaintenance4Properties props4; + vk::PhysicalDeviceSubgroupProperties subgroup_props; + vk::PhysicalDeviceDriverProperties driver_props; + vk::PhysicalDeviceShaderSMBuiltinsPropertiesNV sm_props; + vk::PhysicalDeviceShaderCoreProperties2AMD amd_shader_core_properties2_props; + vk::PhysicalDeviceVulkan12Properties vk12_props; + vk::PhysicalDeviceSubgroupSizeControlPropertiesEXT subgroup_size_control_props; + + props2.pNext = &props3; + props3.pNext = &subgroup_props; + subgroup_props.pNext = &driver_props; + driver_props.pNext = &vk12_props; + + VkBaseOutStructure * last_struct = (VkBaseOutStructure *)&vk12_props; + + if (maintenance4_support) { + last_struct->pNext = (VkBaseOutStructure *)&props4; + last_struct = (VkBaseOutStructure *)&props4; + } + if (sm_builtins) { + last_struct->pNext = (VkBaseOutStructure *)&sm_props; + last_struct = (VkBaseOutStructure *)&sm_props; + } + if (amd_shader_core_properties2) { + last_struct->pNext = (VkBaseOutStructure *)&amd_shader_core_properties2_props; + last_struct = (VkBaseOutStructure *)&amd_shader_core_properties2_props; + } + if (device->subgroup_size_control) { + last_struct->pNext = (VkBaseOutStructure *)&subgroup_size_control_props; + last_struct = (VkBaseOutStructure *)&subgroup_size_control_props; + } + +#if defined(VK_NV_cooperative_matrix2) + vk::PhysicalDeviceCooperativeMatrix2PropertiesNV coopmat2_props; + if (coopmat2_support) { + last_struct->pNext = (VkBaseOutStructure *)&coopmat2_props; + last_struct = (VkBaseOutStructure *)&coopmat2_props; + } +#endif + + device->physical_device.getProperties2(&props2); + device->properties = props2.properties; + + const char* GGML_VK_FORCE_MAX_ALLOCATION_SIZE = getenv("GGML_VK_FORCE_MAX_ALLOCATION_SIZE"); + + if (GGML_VK_FORCE_MAX_ALLOCATION_SIZE != nullptr) { + device->max_memory_allocation_size = std::stoul(GGML_VK_FORCE_MAX_ALLOCATION_SIZE); + } else if (maintenance4_support) { + device->max_memory_allocation_size = std::min(props3.maxMemoryAllocationSize, props4.maxBufferSize); + } else { + device->max_memory_allocation_size = props3.maxMemoryAllocationSize; + } + + device->vendor_id = device->properties.vendorID; + device->subgroup_size = subgroup_props.subgroupSize; + device->uma = device->properties.deviceType == vk::PhysicalDeviceType::eIntegratedGpu; + if (sm_builtins) { + device->shader_core_count = sm_props.shaderSMCount; + } else if (amd_shader_core_properties2) { + device->shader_core_count = amd_shader_core_properties2_props.activeComputeUnitCount; + } else { + device->shader_core_count = 0; + } + device->float_controls_rte_fp16 = vk12_props.shaderRoundingModeRTEFloat16; + + const bool force_disable_f16 = getenv("GGML_VK_DISABLE_F16") != nullptr; + + device->fp16 = !force_disable_f16 && fp16_storage && fp16_compute; + + if (!ggml_vk_khr_cooperative_matrix_support(device->properties, driver_props)) { + device->coopmat_support = false; + } + + std::vector queue_family_props = device->physical_device.getQueueFamilyProperties(); + + // Try to find a non-graphics compute queue and transfer-focused queues + const uint32_t compute_queue_family_index = ggml_vk_find_queue_family_index(queue_family_props, vk::QueueFlagBits::eCompute, vk::QueueFlagBits::eGraphics, -1, 1); + const uint32_t transfer_queue_family_index = ggml_vk_find_queue_family_index(queue_family_props, vk::QueueFlagBits::eTransfer, vk::QueueFlagBits::eCompute | vk::QueueFlagBits::eGraphics, compute_queue_family_index, 1); + + const float priorities[] = { 1.0f, 1.0f }; + device->single_queue = compute_queue_family_index == transfer_queue_family_index && queue_family_props[compute_queue_family_index].queueCount == 1; + + std::vector device_queue_create_infos; + if (compute_queue_family_index != transfer_queue_family_index) { + device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities}); + device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), transfer_queue_family_index, 1, priorities + 1}); + } else if(!device->single_queue) { + device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 2, priorities}); + } else { + device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities}); + } + vk::DeviceCreateInfo device_create_info; + std::vector device_extensions; + vk::PhysicalDeviceFeatures device_features = device->physical_device.getFeatures(); + + VkPhysicalDeviceFeatures2 device_features2; + device_features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2; + device_features2.pNext = nullptr; + device_features2.features = (VkPhysicalDeviceFeatures)device_features; + + VkPhysicalDeviceVulkan11Features vk11_features; + vk11_features.pNext = nullptr; + vk11_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES; + device_features2.pNext = &vk11_features; + + VkPhysicalDeviceVulkan12Features vk12_features; + vk12_features.pNext = nullptr; + vk12_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES; + vk11_features.pNext = &vk12_features; + + last_struct = (VkBaseOutStructure *)&vk12_features; + + VkPhysicalDevicePipelineRobustnessFeaturesEXT pl_robustness_features; + pl_robustness_features.pNext = nullptr; + pl_robustness_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PIPELINE_ROBUSTNESS_FEATURES_EXT; + pl_robustness_features.pipelineRobustness = VK_FALSE; + + if (pipeline_robustness) { + last_struct->pNext = (VkBaseOutStructure *)&pl_robustness_features; + last_struct = (VkBaseOutStructure *)&pl_robustness_features; + device_extensions.push_back("VK_EXT_pipeline_robustness"); + } + + VkPhysicalDeviceSubgroupSizeControlFeaturesEXT subgroup_size_control_features; + subgroup_size_control_features.pNext = nullptr; + subgroup_size_control_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES_EXT; + subgroup_size_control_features.computeFullSubgroups = false; + subgroup_size_control_features.subgroupSizeControl = false; + + if (device->subgroup_size_control) { + last_struct->pNext = (VkBaseOutStructure *)&subgroup_size_control_features; + last_struct = (VkBaseOutStructure *)&subgroup_size_control_features; + } + + VkPhysicalDeviceCooperativeMatrixFeaturesKHR coopmat_features; + coopmat_features.pNext = nullptr; + coopmat_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR; + coopmat_features.cooperativeMatrix = VK_FALSE; + + if (device->coopmat_support) { + last_struct->pNext = (VkBaseOutStructure *)&coopmat_features; + last_struct = (VkBaseOutStructure *)&coopmat_features; + } + +#if defined(VK_NV_cooperative_matrix2) + VkPhysicalDeviceCooperativeMatrix2FeaturesNV coopmat2_features {}; + coopmat2_features.pNext = nullptr; + coopmat2_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_2_FEATURES_NV; + if (coopmat2_support) { + last_struct->pNext = (VkBaseOutStructure *)&coopmat2_features; + last_struct = (VkBaseOutStructure *)&coopmat2_features; + device_extensions.push_back("VK_NV_cooperative_matrix2"); + } +#endif + + vkGetPhysicalDeviceFeatures2(device->physical_device, &device_features2); + + device->fp16 = device->fp16 && vk12_features.shaderFloat16; + + device->pipeline_robustness = pl_robustness_features.pipelineRobustness; + + if (device->subgroup_size_control) { + device->subgroup_min_size = subgroup_size_control_props.minSubgroupSize; + device->subgroup_max_size = subgroup_size_control_props.maxSubgroupSize; + } + + device->subgroup_size_control = device->subgroup_size_control && + (subgroup_size_control_props.requiredSubgroupSizeStages & vk::ShaderStageFlagBits::eCompute) && + subgroup_size_control_features.subgroupSizeControl; + + if (device->subgroup_size_control) { + device->subgroup_require_full_support = subgroup_size_control_features.computeFullSubgroups; + device_extensions.push_back("VK_EXT_subgroup_size_control"); + } + + device->coopmat_support = device->coopmat_support && coopmat_features.cooperativeMatrix; + + if (coopmat2_support) { +#if defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + if (coopmat2_features.cooperativeMatrixWorkgroupScope && + coopmat2_features.cooperativeMatrixFlexibleDimensions && + coopmat2_features.cooperativeMatrixReductions && + coopmat2_features.cooperativeMatrixConversions && + coopmat2_features.cooperativeMatrixPerElementOperations && + coopmat2_features.cooperativeMatrixTensorAddressing && + coopmat2_features.cooperativeMatrixBlockLoads && + vk12_features.bufferDeviceAddress) { + + std::vector flexible_dimensions; + uint32_t count = 0; + + PFN_vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV + _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV = + (PFN_vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV) + vk_instance.instance.getProcAddr("vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV"); + + _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV(device->physical_device, &count, nullptr); + + VkCooperativeMatrixFlexibleDimensionsPropertiesNV empty_prop {}; + empty_prop.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_FLEXIBLE_DIMENSIONS_PROPERTIES_NV; + flexible_dimensions.resize(count, empty_prop); + + _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV(device->physical_device, &count, flexible_dimensions.data()); + + bool found_fp16_128 = false, + found_fp16_256 = false, + found_fp32_128 = false, + found_fp32_256 = false; + // need to support fp16*fp16 with fp16/fp32 accumulator, for workgroupsize 128 + // with 32x16x16 and 256 with 32x32x16. + for (auto &prop : flexible_dimensions) { + if (prop.saturatingAccumulation == VK_FALSE && + prop.scope == VK_SCOPE_WORKGROUP_KHR && + prop.AType == VK_COMPONENT_TYPE_FLOAT16_KHR && + prop.BType == VK_COMPONENT_TYPE_FLOAT16_KHR) { + + if (prop.workgroupInvocations == 128 && + prop.MGranularity <= 32 && + prop.NGranularity <= 16 && + prop.KGranularity <= 16) { + if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR && + prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) { + found_fp16_128 = true; + } + if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && + prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) { + found_fp32_128 = true; + } + } + if (prop.workgroupInvocations == 256 && + prop.MGranularity <= 32 && + prop.NGranularity <= 32 && + prop.KGranularity <= 16) { + if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR && + prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) { + found_fp16_256 = true; + } + if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && + prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) { + found_fp32_256 = true; + } + } + } + } + if (found_fp16_128 && found_fp16_256 && + found_fp32_128 && found_fp32_256 && + coopmat2_props.cooperativeMatrixFlexibleDimensionsMaxDimension >= 512) { + device->coopmat2 = true; + } + } +#endif + } + + if (!vk11_features.storageBuffer16BitAccess) { + std::cerr << "ggml_vulkan: device " << GGML_VK_NAME << idx << " does not support 16-bit storage." << std::endl; + throw std::runtime_error("Unsupported device"); + } + + device_extensions.push_back("VK_KHR_16bit_storage"); + +#ifdef GGML_VULKAN_VALIDATE + device_extensions.push_back("VK_KHR_shader_non_semantic_info"); +#endif + + if (device->fp16) { + device_extensions.push_back("VK_KHR_shader_float16_int8"); + } + + if (device->coopmat_support) { + // Query supported shapes + std::vector cm_props; + + PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR = + (PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR)vkGetInstanceProcAddr(vk_instance.instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR"); + + uint32_t cm_props_num; + + pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR(device->physical_device, &cm_props_num, nullptr); + + cm_props.resize(cm_props_num); + + for (auto& prop : cm_props) { + prop.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR; + } + + pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR(device->physical_device, &cm_props_num, cm_props.data()); + + VK_LOG_DEBUG("ggml_vulkan: Cooperative Matrix Shapes: " << cm_props.size()); + + for (auto& prop : cm_props) { + VK_LOG_DEBUG("ggml_vulkan: M: " << prop.MSize << " N: " << prop.NSize << " K: " << prop.KSize << " A: " << vk::to_string((vk::ComponentTypeKHR)prop.AType) << " B: " << vk::to_string((vk::ComponentTypeKHR)prop.BType) << " C: " << vk::to_string((vk::ComponentTypeKHR)prop.CType) << " Result: " << vk::to_string((vk::ComponentTypeKHR)prop.ResultType) << " saturatingAccumulation: " << prop.saturatingAccumulation << " scope: " << vk::to_string((vk::ScopeKHR)prop.scope)); + + if ((vk::ComponentTypeKHR)prop.AType == vk::ComponentTypeKHR::eFloat16 && + (vk::ComponentTypeKHR)prop.BType == vk::ComponentTypeKHR::eFloat16 && + (vk::ScopeKHR)prop.scope == vk::ScopeKHR::eSubgroup + ) { + if ((vk::ComponentTypeKHR)prop.CType == vk::ComponentTypeKHR::eFloat32 && + (vk::ComponentTypeKHR)prop.ResultType == vk::ComponentTypeKHR::eFloat32) { + // coopmat sizes not set yet + if (device->coopmat_m == 0) { + device->coopmat_acc_f32_support = true; + device->coopmat_m = prop.MSize; + device->coopmat_n = prop.NSize; + device->coopmat_k = prop.KSize; + } else if (device->coopmat_m == prop.MSize && device->coopmat_n == prop.NSize && device->coopmat_k == prop.KSize) { + // Only enable if shape is identical + device->coopmat_acc_f32_support = true; + } + } else if ((vk::ComponentTypeKHR)prop.CType == vk::ComponentTypeKHR::eFloat16 && + (vk::ComponentTypeKHR)prop.ResultType == vk::ComponentTypeKHR::eFloat16) { + // coopmat sizes not set yet + if (device->coopmat_m == 0) { + device->coopmat_acc_f16_support = true; + device->coopmat_m = prop.MSize; + device->coopmat_n = prop.NSize; + device->coopmat_k = prop.KSize; + } else if (device->coopmat_m == prop.MSize && device->coopmat_n == prop.NSize && device->coopmat_k == prop.KSize) { + // Only enable if shape is identical + device->coopmat_acc_f16_support = true; + } + } + } + } + + if (device->coopmat_m == 0 || !device->coopmat_acc_f32_support) { + // No suitable matmul mode found + GGML_LOG_DEBUG("ggml_vulkan: WARNING: No suitable matrix core mode found. Disabling matrix cores.\n"); + device->coopmat_support = false; + } + } + + if (device->coopmat_support) { + device_extensions.push_back("VK_KHR_cooperative_matrix"); + } + + device->name = GGML_VK_NAME + std::to_string(idx); + + device_create_info = { + vk::DeviceCreateFlags(), + device_queue_create_infos, + {}, + device_extensions + }; + device_create_info.setPNext(&device_features2); + device->device = device->physical_device.createDevice(device_create_info); + + // Queues + ggml_vk_create_queue(device, device->compute_queue, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false); + + // Shaders + // Disable matmul tile sizes early if performance low or not supported + switch (device->vendor_id) { +#ifndef GGML_VULKAN_RUN_TESTS + case VK_VENDOR_ID_AMD: + case VK_VENDOR_ID_INTEL: + device->mul_mat_l = false; + device->mul_mat_m = true; + device->mul_mat_s = true; + device->mul_mat_id_l = false; + device->mul_mat_id_m = true; + device->mul_mat_id_s = true; + break; + case VK_VENDOR_ID_APPLE: + device->mul_mat_l = false; + device->mul_mat_m = true; + device->mul_mat_s = false; + device->mul_mat_id_l = false; + device->mul_mat_id_m = true; + device->mul_mat_id_s = false; + break; +#endif + default: + device->mul_mat_l = true; + device->mul_mat_m = true; + device->mul_mat_s = true; + device->mul_mat_id_l = true; + device->mul_mat_id_m = true; + device->mul_mat_id_s = true; + break; + } + + ggml_vk_load_shaders(device); + + if (!device->single_queue) { + const uint32_t transfer_queue_index = compute_queue_family_index == transfer_queue_family_index ? 1 : 0; + ggml_vk_create_queue(device, device->transfer_queue, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true); + } else { + // TODO: Use pointer or reference to avoid copy + device->transfer_queue = device->compute_queue; + } + + device->buffer_type = { + /* .iface = */ ggml_backend_vk_buffer_type_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), idx), + /* .context = */ new ggml_backend_vk_buffer_type_context{ device->name, device }, + }; + + device->fence = device->device.createFence({}); + + device->idx = idx; + + return device; + } + + return vk_instance.devices[idx]; +} + +static void ggml_vk_print_gpu_info(size_t idx) { + GGML_ASSERT(idx < vk_instance.device_indices.size()); + size_t dev_num = vk_instance.device_indices[idx]; + VK_LOG_DEBUG("ggml_vk_print_gpu_info(" << dev_num << ")"); + GGML_ASSERT(vk_instance_initialized); + + std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); + + if (dev_num >= devices.size()) { + std::cerr << "ggml_vulkan: Device with index " << dev_num << " does not exist." << std::endl; + throw std::runtime_error("Device not found"); + } + + vk::PhysicalDevice physical_device = devices[dev_num]; + std::vector ext_props = physical_device.enumerateDeviceExtensionProperties(); + + vk::PhysicalDeviceProperties2 props2; + vk::PhysicalDeviceMaintenance3Properties props3; + vk::PhysicalDeviceSubgroupProperties subgroup_props; + vk::PhysicalDeviceDriverProperties driver_props; + props2.pNext = &props3; + props3.pNext = &subgroup_props; + subgroup_props.pNext = &driver_props; + physical_device.getProperties2(&props2); + + const size_t subgroup_size = subgroup_props.subgroupSize; + const bool uma = props2.properties.deviceType == vk::PhysicalDeviceType::eIntegratedGpu; + + bool fp16_storage = false; + bool fp16_compute = false; + bool coopmat_support = false; + bool coopmat2_support = false; + + for (auto properties : ext_props) { + if (strcmp("VK_KHR_16bit_storage", properties.extensionName) == 0) { + fp16_storage = true; + } else if (strcmp("VK_KHR_shader_float16_int8", properties.extensionName) == 0) { + fp16_compute = true; + } else if (strcmp("VK_KHR_cooperative_matrix", properties.extensionName) == 0 && + !getenv("GGML_VK_DISABLE_COOPMAT")) { + coopmat_support = true; +#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + } else if (strcmp("VK_NV_cooperative_matrix2", properties.extensionName) == 0 && + !getenv("GGML_VK_DISABLE_COOPMAT2")) { + coopmat2_support = true; +#endif + } + } + + if (!ggml_vk_khr_cooperative_matrix_support(props2.properties, driver_props)) { + coopmat_support = false; + } + + const char* GGML_VK_DISABLE_F16 = getenv("GGML_VK_DISABLE_F16"); + bool force_disable_f16 = GGML_VK_DISABLE_F16 != nullptr; + + bool fp16 = !force_disable_f16 && fp16_storage && fp16_compute; + + vk::PhysicalDeviceFeatures device_features = physical_device.getFeatures(); + + VkPhysicalDeviceFeatures2 device_features2; + device_features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2; + device_features2.pNext = nullptr; + device_features2.features = (VkPhysicalDeviceFeatures)device_features; + + VkPhysicalDeviceVulkan11Features vk11_features; + vk11_features.pNext = nullptr; + vk11_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES; + device_features2.pNext = &vk11_features; + + VkPhysicalDeviceVulkan12Features vk12_features; + vk12_features.pNext = nullptr; + vk12_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES; + vk11_features.pNext = &vk12_features; + + // Pointer to the last chain element + VkBaseOutStructure * last_struct = (VkBaseOutStructure *)&vk12_features; + + VkPhysicalDeviceCooperativeMatrixFeaturesKHR coopmat_features; + coopmat_features.pNext = nullptr; + coopmat_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR; + coopmat_features.cooperativeMatrix = VK_FALSE; + + if (coopmat_support) { + last_struct->pNext = (VkBaseOutStructure *)&coopmat_features; + last_struct = (VkBaseOutStructure *)&coopmat_features; + } + + vkGetPhysicalDeviceFeatures2(physical_device, &device_features2); + + fp16 = fp16 && vk12_features.shaderFloat16; + + coopmat_support = coopmat_support && coopmat_features.cooperativeMatrix; + + std::string matrix_cores = coopmat2_support ? "NV_coopmat2" : coopmat_support ? "KHR_coopmat" : "none"; + + std::string device_name = props2.properties.deviceName.data(); + GGML_LOG_DEBUG("ggml_vulkan: %zu = %s (%s) | uma: %d | fp16: %d | warp size: %zu | matrix cores: %s\n", + idx, device_name.c_str(), driver_props.driverName.data(), uma, fp16, subgroup_size, matrix_cores.c_str()); + + if (props2.properties.deviceType == vk::PhysicalDeviceType::eCpu) { + GGML_LOG_DEBUG("ggml_vulkan: Warning: Device type is CPU. This is probably not the device you want.\n"); + } +} + +static bool ggml_vk_instance_validation_ext_available(const std::vector& instance_extensions); +static bool ggml_vk_instance_portability_enumeration_ext_available(const std::vector& instance_extensions); + +void ggml_vk_instance_init() { + if (vk_instance_initialized) { + return; + } + VK_LOG_DEBUG("ggml_vk_instance_init()"); + + vk_instance_initialized = true; + + vk::ApplicationInfo app_info{ "ggml-vulkan", 1, nullptr, 0, VK_API_VERSION }; + + const std::vector instance_extensions = vk::enumerateInstanceExtensionProperties(); + const bool validation_ext = ggml_vk_instance_validation_ext_available(instance_extensions); +#ifdef __APPLE__ + const bool portability_enumeration_ext = ggml_vk_instance_portability_enumeration_ext_available(instance_extensions); +#endif + + std::vector layers; + + if (validation_ext) { + layers.push_back("VK_LAYER_KHRONOS_validation"); + } + std::vector extensions; + if (validation_ext) { + extensions.push_back("VK_EXT_validation_features"); + } +#ifdef __APPLE__ + if (portability_enumeration_ext) { + extensions.push_back("VK_KHR_portability_enumeration"); + } +#endif + vk::InstanceCreateInfo instance_create_info(vk::InstanceCreateFlags{}, &app_info, layers, extensions); +#ifdef __APPLE__ + if (portability_enumeration_ext) { + instance_create_info.flags |= vk::InstanceCreateFlagBits::eEnumeratePortabilityKHR; + } +#endif + + std::vector features_enable; + vk::ValidationFeaturesEXT validation_features; + + if (validation_ext) { + features_enable = { vk::ValidationFeatureEnableEXT::eBestPractices }; + validation_features = { + features_enable, + {}, + }; + validation_features.setPNext(nullptr); + instance_create_info.setPNext(&validation_features); + GGML_LOG_DEBUG("ggml_vulkan: Validation layers enabled\n"); + } + vk_instance.instance = vk::createInstance(instance_create_info); + + size_t num_available_devices = vk_instance.instance.enumeratePhysicalDevices().size(); + + // Emulate behavior of CUDA_VISIBLE_DEVICES for Vulkan + char * devices_env = getenv("GGML_VK_VISIBLE_DEVICES"); + if (devices_env != nullptr) { + std::string devices(devices_env); + std::replace(devices.begin(), devices.end(), ',', ' '); + + std::stringstream ss(devices); + size_t tmp; + while (ss >> tmp) { + if(tmp >= num_available_devices) { + std::cerr << "ggml_vulkan: Invalid device index " << tmp << " in GGML_VK_VISIBLE_DEVICES." << std::endl; + throw std::runtime_error("Invalid Vulkan device index"); + } + vk_instance.device_indices.push_back(tmp); + } + } else { + std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); + + // Make sure at least one device exists + if (devices.empty()) { + std::cerr << "ggml_vulkan: Error: No devices found." << std::endl; + GGML_ABORT("fatal error"); + } + + // Default to using all dedicated GPUs + for (size_t i = 0; i < devices.size(); i++) { + vk::PhysicalDeviceProperties2 new_props; + vk::PhysicalDeviceDriverProperties new_driver; + vk::PhysicalDeviceIDProperties new_id; + new_props.pNext = &new_driver; + new_driver.pNext = &new_id; + devices[i].getProperties2(&new_props); + + if (new_props.properties.deviceType == vk::PhysicalDeviceType::eDiscreteGpu) { + // Check if there are two physical devices corresponding to the same GPU + auto old_device = std::find_if( + vk_instance.device_indices.begin(), + vk_instance.device_indices.end(), + [&devices, &new_id](const size_t k){ + vk::PhysicalDeviceProperties2 old_props; + vk::PhysicalDeviceIDProperties old_id; + old_props.pNext = &old_id; + devices[k].getProperties2(&old_props); + return std::equal(std::begin(old_id.deviceUUID), std::end(old_id.deviceUUID), std::begin(new_id.deviceUUID)); + } + ); + if (old_device == vk_instance.device_indices.end()) { + vk_instance.device_indices.push_back(i); + } else { + // There can be two physical devices corresponding to the same GPU if there are 2 different drivers + // This can cause error when splitting layers aross the devices, need to keep only 1 + VK_LOG_DEBUG("Device " << i << " and device " << *old_device << " have the same deviceUUID"); + + vk::PhysicalDeviceProperties2 old_props; + vk::PhysicalDeviceDriverProperties old_driver; + old_props.pNext = &old_driver; + devices[*old_device].getProperties2(&old_props); + + std::map driver_priorities {}; + int old_priority = std::numeric_limits::max(); + int new_priority = std::numeric_limits::max(); + + // Check https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/VkDriverId.html for the list of driver id + // Smaller number -> higher priority + switch (old_props.properties.vendorID) { + case VK_VENDOR_ID_AMD: + driver_priorities[vk::DriverId::eMesaRadv] = 1; + driver_priorities[vk::DriverId::eAmdOpenSource] = 2; + driver_priorities[vk::DriverId::eAmdProprietary] = 3; + break; + case VK_VENDOR_ID_INTEL: + driver_priorities[vk::DriverId::eIntelOpenSourceMESA] = 1; + driver_priorities[vk::DriverId::eIntelProprietaryWindows] = 2; + break; + case VK_VENDOR_ID_NVIDIA: + driver_priorities[vk::DriverId::eNvidiaProprietary] = 1; +#if defined(VK_API_VERSION_1_3) && VK_HEADER_VERSION >= 235 + driver_priorities[vk::DriverId::eMesaNvk] = 2; +#endif + break; + } + + if (driver_priorities.count(old_driver.driverID)) { + old_priority = driver_priorities[old_driver.driverID]; + } + if (driver_priorities.count(new_driver.driverID)) { + new_priority = driver_priorities[new_driver.driverID]; + } + + if (new_priority < old_priority) { + auto r = std::remove(vk_instance.device_indices.begin(), vk_instance.device_indices.end(), *old_device); + vk_instance.device_indices.erase(r, vk_instance.device_indices.end()); + vk_instance.device_indices.push_back(i); + + VK_LOG_DEBUG("Prioritize device " << i << " driver " << new_driver.driverName << " over device " << *old_device << " driver " << old_driver.driverName); + } + else { + VK_LOG_DEBUG("Prioritize device " << *old_device << " driver " << old_driver.driverName << " over device " << i << " driver " << new_driver.driverName << std::endl); + } + } + } + } + + // If no dedicated GPUs found, fall back to GPU 0 + if (vk_instance.device_indices.empty()) { + vk_instance.device_indices.push_back(0); + } + } + GGML_LOG_DEBUG("ggml_vulkan: Found %zu Vulkan devices:\n", vk_instance.device_indices.size()); + + for (size_t i = 0; i < vk_instance.device_indices.size(); i++) { + ggml_vk_print_gpu_info(i); + } +} + +static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) { + VK_LOG_DEBUG("ggml_vk_init(" << ctx->name << ", " << idx << ")"); + ggml_vk_instance_init(); + GGML_ASSERT(idx < vk_instance.device_indices.size()); + + ctx->name = GGML_VK_NAME + std::to_string(idx); + + ctx->device = ggml_vk_get_device(idx); + + ctx->semaphore_idx = 0; + ctx->event_idx = 0; + + ctx->prealloc_size_x = 0; + ctx->prealloc_size_y = 0; + ctx->prealloc_size_split_k = 0; + + ctx->fence = ctx->device->device.createFence({}); + +#ifdef GGML_VULKAN_CHECK_RESULTS + const char* skip_checks = getenv("GGML_VULKAN_SKIP_CHECKS"); + vk_skip_checks = (skip_checks == NULL ? 0 : atoi(skip_checks)); + const char* output_tensor = getenv("GGML_VULKAN_OUTPUT_TENSOR"); + vk_output_tensor = (output_tensor == NULL ? 0 : atoi(output_tensor)); +#endif +} + +static vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type type) { + VK_LOG_DEBUG("ggml_vk_get_to_fp16()"); + switch (type) { + case GGML_TYPE_F32: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return nullptr; + } + + return ctx->device->pipeline_dequant[type]; +} + +static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec) { + VK_LOG_DEBUG("ggml_vk_get_mul_mat_mat_pipeline(" << ggml_type_name(src0_type) << ", " << ggml_type_name(src1_type) << ")"); + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_f32; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_f32_f16; + } + if (prec == GGML_PREC_DEFAULT && ctx->device->fp16 && !(ctx->device->coopmat_support && !ctx->device->coopmat_acc_f16_support)) { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_f16_f32.f16acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_f16.f16acc; + } + } else { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_f16_f32.f32acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_f16.f32acc; + } + } + + if (src1_type != GGML_TYPE_F32 && !ctx->device->coopmat2) { + return nullptr; + } + + switch (src0_type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return nullptr; + } + + if (ctx->device->coopmat2) { + assert(src1_type == GGML_TYPE_F16); + return ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f16acc; + } + return ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f32acc; +} + +static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * ctx, ggml_type a_type, ggml_type b_type, uint32_t num_cols) { + VK_LOG_DEBUG("ggml_vk_get_dequantize_mul_mat_vec()"); + GGML_ASSERT(b_type == GGML_TYPE_F32 || b_type == GGML_TYPE_F16); + GGML_ASSERT(num_cols >= 1 && num_cols <= mul_mat_vec_max_cols); + + switch (a_type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return nullptr; + } + + return b_type == GGML_TYPE_F32 ? ctx->device->pipeline_dequant_mul_mat_vec_f32_f32[a_type][num_cols-1] : ctx->device->pipeline_dequant_mul_mat_vec_f16_f32[a_type][num_cols-1]; +} + +static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec) { + VK_LOG_DEBUG("ggml_vk_get_mul_mat_mat_id_pipeline()"); + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_id_f32; + } + if (prec == GGML_PREC_DEFAULT && ctx->device->fp16 && !(ctx->device->coopmat_support && !ctx->device->coopmat_acc_f16_support)) { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_id_f16_f32.f16acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_id_f16.f16acc; + } + } else { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_id_f16_f32.f32acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_id_f16.f32acc; + } + } + + GGML_ASSERT(src1_type == GGML_TYPE_F32); + + switch (src0_type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return nullptr; + } + + return ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat_id[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat_id[src0_type].f32acc; +} + +static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context * ctx, ggml_type a_type, ggml_type b_type) { + VK_LOG_DEBUG("ggml_vk_get_dequantize_mul_mat_vec()"); + GGML_ASSERT(b_type == GGML_TYPE_F32); + + switch (a_type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return nullptr; + } + + return ctx->device->pipeline_dequant_mul_mat_vec_id_f32[a_type]; +} + +static vk_buffer ggml_vk_pool_malloc(ggml_backend_vk_context * ctx, size_t size) { + VK_LOG_DEBUG("ggml_vk_pool_malloc(" << size << ")"); + VK_LOG_MEMORY("ggml_vk_pool_malloc"); + + int best_i = -1; + size_t best_size = std::numeric_limits::max(); //smallest unused buffer that fits our needs + int worst_i = -1; + size_t worst_size = 0; //largest unused buffer seen so far + for (int i = 0; i < MAX_VK_BUFFERS; ++i) { + vk_buffer &b = ctx->buffer_pool[i]; + if (b != nullptr && b->size >= size && b->size < best_size) { + best_i = i; + best_size = b->size; + } + if (b != nullptr && b->size > worst_size) { + worst_i = i; + worst_size = b->size; + } + } + if(best_i != -1) { + //found the smallest buffer that fits our needs + vk_buffer b = ctx->buffer_pool[best_i]; + ctx->buffer_pool[best_i].reset(); + return b; + } + if(worst_i != -1) { + //no buffer that fits our needs, resize largest one to save memory + vk_buffer& b = ctx->buffer_pool[worst_i]; + ggml_vk_destroy_buffer(b); + } + + return ggml_vk_create_buffer_device(ctx->device, size); +} + +static void ggml_vk_pool_free(ggml_backend_vk_context * ctx, vk_buffer& buffer) { + VK_LOG_DEBUG("ggml_vk_pool_free(" << buffer->size << ")"); + for (int i = 0; i < MAX_VK_BUFFERS; ++i) { + vk_buffer& b = ctx->buffer_pool[i]; + if (b == nullptr) { + b = buffer; + return; + } + } + std::cerr << "ggml_vulkan: WARNING: vk buffer pool full, increase MAX_VK_BUFFERS" << std::endl; + ggml_vk_destroy_buffer(buffer); +} + +// Returns an available temporary buffer that may only be used temporarily, it will be reused +static vk_buffer ggml_vk_create_buffer_temp(ggml_backend_vk_context * ctx, size_t size) { + // Try to find existing temp buffer with enough capacity + for (auto& buffer : ctx->gc.temp_buffers) { + if (buffer->size >= size) { + return buffer; + } + } + + VK_LOG_MEMORY("ggml_vk_create_buffer_temp(" << size << ")"); + + // Otherwise create new buffer + vk_buffer buf = ggml_vk_pool_malloc(ctx, size); + ctx->gc.temp_buffers.push_back(buf); + + return buf; +} + +static void * ggml_vk_host_malloc(vk_device& device, size_t size) { + VK_LOG_MEMORY("ggml_vk_host_malloc(" << size << ")"); + vk_buffer buf = ggml_vk_create_buffer(device, size, + vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached, + vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); + + if(!(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible)) { + fprintf(stderr, "WARNING: failed to allocate %.2f MB of pinned memory\n", + size/1024.0/1024.0); + device->device.freeMemory(buf->device_memory); + device->device.destroyBuffer(buf->buffer); + return nullptr; + } + + device->pinned_memory.push_back(std::make_tuple(buf->ptr, size, buf)); + + return buf->ptr; +} + +static void ggml_vk_host_free(vk_device& device, void* ptr) { + if (ptr == nullptr) { + return; + } + VK_LOG_MEMORY("ggml_vk_host_free(" << ptr << ")"); + vk_buffer buf; + size_t index; + for (size_t i = 0; i < device->pinned_memory.size(); i++) { + const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]); + const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]); + if (ptr >= addr && ptr < endr) { + buf = std::get<2>(device->pinned_memory[i]); + index = i; + break; + } + } + if (buf == nullptr) { + fprintf(stderr, "WARNING: failed to free pinned memory: memory not in map\n"); + return; + } + + ggml_vk_destroy_buffer(buf); + + device->pinned_memory.erase(device->pinned_memory.begin() + index); +} + +static void ggml_vk_host_get(vk_device& device, const void * ptr, vk_buffer& buf, size_t& buf_offset) { + buf = nullptr; + buf_offset = 0; + for (size_t i = 0; i < device->pinned_memory.size(); i++) { + const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]); + const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]); + if (ptr >= addr && ptr < endr) { + buf = std::get<2>(device->pinned_memory[i]); + buf_offset = ((const uint8_t *)ptr) - addr; + break; + } + } +} + +static vk_submission ggml_vk_begin_submission(vk_device& device, vk_queue& q, bool one_time = true) { + vk_submission s; + s.buffer = ggml_vk_create_cmd_buffer(device, q); + if (one_time) { + s.buffer.begin({ vk::CommandBufferUsageFlagBits::eOneTimeSubmit }); + } else { + s.buffer.begin({ vk::CommandBufferUsageFlags{} }); + } + + return s; +} + + + +static void ggml_vk_dispatch_pipeline(ggml_backend_vk_context* ctx, vk_context& subctx, vk_pipeline& pipeline, std::initializer_list const& descriptor_buffer_infos, size_t push_constant_size, const void* push_constants, std::array elements) { + const uint32_t wg0 = CEIL_DIV(elements[0], pipeline->wg_denoms[0]); + const uint32_t wg1 = CEIL_DIV(elements[1], pipeline->wg_denoms[1]); + const uint32_t wg2 = CEIL_DIV(elements[2], pipeline->wg_denoms[2]); + VK_LOG_DEBUG("ggml_vk_dispatch_pipeline(" << pipeline->name << ", {"; + for (auto& buffer : descriptor_buffer_infos) { + std::cerr << "(" << buffer.buffer << ", " << buffer.offset << ", " << buffer.range << "), "; + } + std::cerr << "}, (" << wg0 << "," << wg1 << "," << wg2 << "))"); + GGML_ASSERT(pipeline->descriptor_set_idx < pipeline->descriptor_sets.size()); + GGML_ASSERT(descriptor_buffer_infos.size() == pipeline->parameter_count); + + vk::DescriptorSet& descriptor_set = pipeline->descriptor_sets[pipeline->descriptor_set_idx++]; + vk::WriteDescriptorSet write_descriptor_set{ descriptor_set, 0, 0, pipeline->parameter_count, vk::DescriptorType::eStorageBuffer, nullptr, descriptor_buffer_infos.begin() }; + ctx->device->device.updateDescriptorSets({ write_descriptor_set }, {}); + + subctx->s->buffer.pushConstants(pipeline->layout, vk::ShaderStageFlagBits::eCompute, 0, push_constant_size, push_constants); + subctx->s->buffer.bindPipeline(vk::PipelineBindPoint::eCompute, pipeline->pipeline); + subctx->s->buffer.bindDescriptorSets(vk::PipelineBindPoint::eCompute, + pipeline->layout, + 0, + { descriptor_set }, + {}); + subctx->s->buffer.dispatch(wg0, wg1, wg2); +} + +static void ggml_vk_end_submission(vk_submission& s, std::vector wait_semaphores, std::vector signal_semaphores) { + s.buffer.end(); + + s.wait_semaphores = std::move(wait_semaphores); + s.signal_semaphores = std::move(signal_semaphores); +} + +static void ggml_vk_ctx_end(vk_context& ctx) { + VK_LOG_DEBUG("ggml_vk_ctx_end(" << ctx << ", " << ctx->seqs.size() << ")"); + if (ctx->s == nullptr) { + return; + } + + ctx->s->buffer.end(); + ctx->s = nullptr; +} + +static void ggml_vk_ctx_begin(vk_device& device, vk_context& subctx) { + VK_LOG_DEBUG("ggml_vk_ctx_begin(" << device->name << ")"); + if (subctx->s != nullptr) { + ggml_vk_ctx_end(subctx); + } + + subctx->seqs.push_back({ ggml_vk_begin_submission(device, *subctx->q) }); + subctx->s = subctx->seqs[subctx->seqs.size() - 1].data(); +} + +static size_t ggml_vk_align_size(size_t width, size_t align) { + VK_LOG_DEBUG("ggml_vk_align_size(" << width << ", " << align << ")"); + return CEIL_DIV(width, align) * align; +} + +static void deferred_memcpy(void * dst, const void * src, size_t size, std::vector* memcpys = nullptr) { + if (memcpys == nullptr) { + memcpy(dst, src, size); + } else { + memcpys->emplace_back(dst, src, size); + } +} + +static void ggml_vk_ensure_sync_staging_buffer(vk_device& device, size_t size) { + if (device->sync_staging == nullptr || device->sync_staging->size < size) { + VK_LOG_MEMORY("ggml_vk_ensure_sync_staging_buffer(" << size << ")"); + ggml_vk_destroy_buffer(device->sync_staging); + device->sync_staging = ggml_vk_create_buffer_check(device, size, + vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached, + vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); + } +} + +static void ggml_vk_buffer_write_nc_async(ggml_backend_vk_context * ctx, vk_context& subctx, vk_buffer& dst, size_t offset, const ggml_tensor * tensor, bool sync_staging = false) { + VK_LOG_DEBUG("ggml_vk_buffer_write_nc_async(" << tensor << ")"); + GGML_ASSERT(!ggml_is_contiguous(tensor)); + // Buffer is already mapped + if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { + std::cerr << "ggml_vulkan: buffer_write_nc_async dst buffer is host_visible. Use synchronous write." << std::endl; + GGML_ABORT("fatal error"); + } + // Check if src is pinned memory + vk_buffer buf = nullptr; + size_t buf_offset = 0; + ggml_vk_host_get(ctx->device, tensor->data, buf, buf_offset); + + const uint64_t ne0 = tensor->ne[0]; + const uint64_t ne1 = tensor->ne[1]; + const uint64_t ne2 = tensor->ne[2]; + const uint64_t ne3 = tensor->ne[3]; + const uint64_t nb0 = tensor->nb[0]; + const uint64_t nb1 = tensor->nb[1]; + const uint64_t nb2 = tensor->nb[2]; + const uint64_t nb3 = tensor->nb[3]; + const ggml_type type = tensor->type; + const uint64_t ts = ggml_type_size(type); + const uint64_t bs = ggml_blck_size(type); + + const uint64_t dstnb0 = ts; + const uint64_t dstnb1 = dstnb0*(ne0/bs); + const uint64_t dstnb2 = dstnb1*ne1; + const uint64_t dstnb3 = dstnb2*ne2; + + const uint64_t ne = ggml_nelements(tensor); + + if (buf != nullptr) { + // Memory is pinned, use as staging buffer + std::vector slices; + + for (uint64_t i3 = 0; i3 < ne3; i3++) { + for (uint64_t i2 = 0; i2 < ne2; i2++) { + // Find longest contiguous slice + if (ne1*nb1 == dstnb2) { + slices.push_back({ buf_offset + i3*nb3 + i2*nb2, offset + i3*dstnb3 + i2*dstnb2, dstnb2 }); + } else { + for (uint64_t i1 = 0; i1 < ne1; i1++) { + if (ne0*nb0/bs == dstnb1) { + slices.push_back({ buf_offset + i3*nb3 + i2*nb2 + i1*nb1, offset + i3*dstnb3 + i2*dstnb2 + i1*dstnb1, dstnb1 }); + } else { + const uint64_t s_off = buf_offset + i3*nb3 + i2*nb2 + i1*nb1; + const uint64_t d_off = offset + i3*dstnb3 + i2*dstnb2 + i1*dstnb1; + for (uint64_t i0 = 0; i0 < ne0; i0++) { + slices.push_back({ s_off + i1*nb0, d_off + i0*dstnb0, dstnb0 }); + } + } + } + } + } + } + + ggml_vk_sync_buffers(subctx); + subctx->s->buffer.copyBuffer(buf->buffer, dst->buffer, slices); + return; + } + + if (!sync_staging) { + GGML_ABORT("Asynchronous write to non-pinned memory not supported"); + } + + // Staging buffer required + vk_buffer& staging = ctx->device->sync_staging; + const uint64_t copy_size = ts*ne/bs; + ggml_vk_ensure_sync_staging_buffer(ctx->device, copy_size); + VkBufferCopy buf_copy{ 0, offset, copy_size }; + + ggml_vk_sync_buffers(subctx); + vkCmdCopyBuffer(subctx->s->buffer, (VkBuffer)staging->buffer, (VkBuffer)dst->buffer, 1, &buf_copy); + + for (uint64_t i3 = 0; i3 < ne3; i3++) { + for (uint64_t i2 = 0; i2 < ne2; i2++) { + // Find longest contiguous slice + if (ne1*nb1 == dstnb2) { + deferred_memcpy((uint8_t *)staging->ptr + i3*dstnb3 + i2*dstnb2, (const uint8_t *) tensor->data + buf_offset + i3*nb3 + i2*nb2, dstnb2, &subctx->in_memcpys); + } else { + for (uint64_t i1 = 0; i1 < ne1; i1++) { + if (ne0*nb0/bs == dstnb1) { + deferred_memcpy((uint8_t *)staging->ptr + i3*dstnb3 + i2*dstnb2 + i1*dstnb1, (const uint8_t *) tensor->data + buf_offset + i3*nb3 + i2*nb2 + i1*nb1, dstnb1, &subctx->in_memcpys); + } else { + const uint64_t s_off = buf_offset + i3*nb3 + i2*nb2 + i1*nb1; + const uint64_t d_off = i3*dstnb3 + i2*dstnb2 + i1*dstnb1; + for (uint64_t i0 = 0; i0 < ne0; i0++) { + deferred_memcpy((uint8_t *)staging->ptr + d_off + i0*dstnb0, (const uint8_t *) tensor->data + s_off + i0*nb0, dstnb0, &subctx->in_memcpys); + } + } + } + } + } + } +} + +static void ggml_vk_buffer_write_2d_async(vk_context subctx, vk_buffer& dst, size_t offset, const void * src, size_t spitch, size_t width, size_t height, bool sync_staging = false) { + VK_LOG_DEBUG("ggml_vk_buffer_write_2d_async(" << width << ", " << height << ")"); + // Buffer is already mapped + if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { + std::cerr << "ggml_vulkan: buffer_write_async dst buffer is host_visible. Use synchronous write." << std::endl; + GGML_ABORT("fatal error"); + } + // Check if src is pinned memory + vk_buffer buf = nullptr; + size_t buf_offset = 0; + ggml_vk_host_get(dst->device, src, buf, buf_offset); + + if (buf != nullptr) { + // Memory is pinned, use as staging buffer + std::vector slices(1); + if (width == spitch) { + // Only do single write if stride is equal + slices[0].srcOffset = buf_offset; + slices[0].dstOffset = offset; + slices[0].size = width * height; + } else { + slices.resize(height); + for (size_t i = 0; i < height; i++) { + slices[i].srcOffset = buf_offset + i * spitch; + slices[i].dstOffset = offset + i * width; + slices[i].size = width; + } + } + + ggml_vk_sync_buffers(subctx); + subctx->s->buffer.copyBuffer(buf->buffer, dst->buffer, slices); + return; + } + VK_LOG_DEBUG("STAGING"); + + if (!sync_staging) { + GGML_ABORT("Asynchronous write to non-pinned memory not supported"); + } + + // Staging buffer required + const size_t copy_size = width*height; + ggml_vk_ensure_sync_staging_buffer(dst->device, copy_size); + + vk_buffer& staging_buffer = dst->device->sync_staging; + + VkBufferCopy buf_copy = { + 0, + offset, + copy_size}; + + ggml_vk_sync_buffers(subctx); + vkCmdCopyBuffer(subctx->s->buffer, (VkBuffer)staging_buffer->buffer, (VkBuffer)dst->buffer, 1, &buf_copy); + + if (width == spitch) { + deferred_memcpy((uint8_t *)staging_buffer->ptr, src, width * height, &subctx->in_memcpys); + } else { + for (size_t i = 0; i < height; i++) { + deferred_memcpy((uint8_t *)staging_buffer->ptr + i * width, (const uint8_t *) src + i * spitch, width, &subctx->in_memcpys); + } + } +} + +static void ggml_vk_buffer_write_async(vk_context subctx, vk_buffer& dst, size_t offset, const void * src, size_t size, bool sync_staging = false) { + VK_LOG_DEBUG("ggml_vk_buffer_write_async(" << size << ")"); + return ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, size, size, 1, sync_staging); +} + +static void ggml_vk_buffer_write_2d(vk_buffer& dst, size_t offset, const void * src, size_t spitch, size_t width, size_t height) { + VK_LOG_DEBUG("ggml_vk_buffer_write_2d(" << width << ", " << height << ")"); + // Buffer is already mapped + if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { + GGML_ASSERT(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent); + + for (size_t i = 0; i < height; i++) { + memcpy((uint8_t *)dst->ptr + offset + i * width, (const uint8_t *) src + i * spitch, width); + } + } else { + vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue); + ggml_vk_ctx_begin(dst->device, subctx); + ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, spitch, width, height, true); + ggml_vk_ctx_end(subctx); + + for (auto& cpy : subctx->in_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + + ggml_vk_submit(subctx, dst->device->fence); + VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_buffer_write_2d waitForFences"); + dst->device->device.resetFences({ dst->device->fence }); + } +} + +static void ggml_vk_buffer_write(vk_buffer& dst, size_t offset, const void * src, size_t size) { + VK_LOG_DEBUG("ggml_vk_buffer_write(" << size << ")"); + ggml_vk_buffer_write_2d(dst, offset, src, 0, size, 1); +} + +static void ggml_vk_buffer_read_2d_async(vk_context subctx, vk_buffer& src, size_t offset, void * dst, size_t spitch, size_t dpitch, size_t width, size_t height, bool sync_staging = false) { + VK_LOG_DEBUG("ggml_vk_buffer_read_2d_async(offset=" << offset << ", width=" << width << ", height=" << height << ")"); + GGML_ASSERT(width > 0); + GGML_ASSERT(height > 0); + GGML_ASSERT(src != nullptr); + + // TODO: staging_offset is not used + + // Check if dst is pinned memory + vk_buffer buf = nullptr; + size_t buf_offset = 0; + ggml_vk_host_get(src->device, dst, buf, buf_offset); + + std::vector slices(1); + if (width == spitch && width == dpitch) { + // Only do single write if stride is equal + slices[0].srcOffset = offset; + slices[0].dstOffset = buf_offset; + slices[0].size = width * height; + } else { + slices.resize(height); + for (size_t i = 0; i < height; i++) { + slices[i].srcOffset = offset + i * spitch; + slices[i].dstOffset = buf_offset + i * dpitch; + slices[i].size = width; + } + } + + if (buf != nullptr) { + // Memory is pinned, use as staging buffer + ggml_vk_sync_buffers(subctx); + subctx->s->buffer.copyBuffer(src->buffer, buf->buffer, slices); + + return; + } + VK_LOG_DEBUG("STAGING"); + + if (!sync_staging) { + GGML_ABORT("Asynchronous read from non-pinned memory not supported"); + } + + // Fall back to staging buffer + const size_t copy_size = dpitch * height; + ggml_vk_ensure_sync_staging_buffer(src->device, copy_size); + + vk_buffer& staging_buffer = src->device->sync_staging; + + ggml_vk_sync_buffers(subctx); + subctx->s->buffer.copyBuffer(src->buffer, staging_buffer->buffer, slices); + + deferred_memcpy(dst, staging_buffer->ptr, copy_size, &subctx->out_memcpys); +} + +static void ggml_vk_buffer_read_async(vk_context subctx, vk_buffer& src, size_t offset, void * dst, size_t size, bool sync_staging = false) { + return ggml_vk_buffer_read_2d_async(subctx, src, offset, dst, size, size, size, 1, sync_staging); +} + +static void ggml_vk_buffer_read(vk_buffer& src, size_t offset, void * dst, size_t size) { + VK_LOG_DEBUG("ggml_vk_buffer_read(" << src->buffer << ", " << offset << ", " << size << ")"); + + // If the device is not an UMA device the memory is host-accessible through rebar. While writing + // through PCIe is sufficient fast reading back data from PCIe is slower than going through + // the HW device to host copy path. + if(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible && src->device->uma) { + GGML_ASSERT(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent); + + memcpy(dst, (uint8_t *) src->ptr + offset, size); + } else { + vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue); + ggml_vk_ctx_begin(src->device, subctx); + ggml_vk_buffer_read_async(subctx, src, offset, dst, size, true); + ggml_vk_ctx_end(subctx); + + ggml_vk_submit(subctx, src->device->fence); + VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_read waitForFences"); + src->device->device.resetFences({ src->device->fence }); + + for (auto& cpy : subctx->out_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + } +} + +static void ggml_vk_buffer_copy_async(vk_context& ctx, vk_buffer& dst, size_t dst_offset, vk_buffer& src, size_t src_offset, size_t size) { + VK_LOG_DEBUG("ggml_vk_buffer_copy_async(" << size << ")"); + // Make sure both buffers are on same device + GGML_ASSERT(src->device == dst->device); + + VkBufferCopy bc{ src_offset, dst_offset, size }; + + vkCmdCopyBuffer(ctx->s->buffer, (VkBuffer)src->buffer, (VkBuffer)dst->buffer, 1, &bc); +} + +static void ggml_vk_buffer_copy(vk_buffer& dst, size_t dst_offset, vk_buffer& src, size_t src_offset, size_t size) { + if (src->device == dst->device) { + VK_LOG_DEBUG("ggml_vk_buffer_copy(SINGLE_DEVICE, " << size << ")"); + // Copy within the device + vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue); + ggml_vk_ctx_begin(src->device, subctx); + ggml_vk_buffer_copy_async(subctx, dst, dst_offset, src, src_offset, size); + ggml_vk_ctx_end(subctx); + ggml_vk_submit(subctx, src->device->fence); + VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_copy waitForFences"); + src->device->device.resetFences({ src->device->fence }); + } else { + VK_LOG_DEBUG("ggml_vk_buffer_copy(MULTI_DEVICE, " << size << ")"); + // Copy device to device + ggml_vk_ensure_sync_staging_buffer(src->device, size); + ggml_vk_ensure_sync_staging_buffer(dst->device, size); + + // Copy to src staging buffer + ggml_vk_buffer_copy(src->device->sync_staging, 0, src, src_offset, size); + // memcpy to dst staging buffer + memcpy(dst->device->sync_staging->ptr, src->device->sync_staging->ptr, size); + // Copy to dst buffer + ggml_vk_buffer_copy(dst, dst_offset, dst->device->sync_staging, 0, size); + } +} + +static void ggml_vk_buffer_memset(vk_buffer& dst, size_t offset, uint32_t c, size_t size) { + VK_LOG_DEBUG("ggml_vk_buffer_memset(" << offset << ", " << c << ", " << size << ")"); + + vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue); + ggml_vk_ctx_begin(dst->device, subctx); + subctx->s->buffer.fillBuffer(dst->buffer, offset, size, c); + ggml_vk_ctx_end(subctx); + + ggml_vk_submit(subctx, dst->device->fence); + VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_memset waitForFences"); + dst->device->device.resetFences({ dst->device->fence }); +} + +static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, int m, int n, int k, const vk_pipeline& pipeline) { + VK_LOG_DEBUG("ggml_vk_guess_split_k(" << m << ", " << n << ", " << k << ")"); + + uint32_t split_k = 1; + if (ctx->device->shader_core_count != 0 && m >= (int)pipeline->wg_denoms[0] && n >= (int)pipeline->wg_denoms[1]) { + // If k is 'large' and the SMs will fill less than halfway, use split_k. + uint32_t m_tiles = CEIL_DIV(m, pipeline->wg_denoms[0]); + uint32_t n_tiles = CEIL_DIV(n, pipeline->wg_denoms[1]); + if (k >= 2048 && m_tiles * n_tiles < ctx->device->shader_core_count / 2) { + split_k = ctx->device->shader_core_count / (m_tiles * n_tiles); + // Clamp to 2 or 4 + split_k = std::min(split_k, 4u); + if (split_k == 3) { + split_k = 2; + } + } + } + + return split_k; +} + +static vk_pipeline ggml_vk_guess_matmul_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, bool aligned) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline(" << m << ", " << n << ", " << aligned << ")"); + + if (ctx->device->coopmat2) { + if ((ctx->device->mul_mat_l && (m % mmp->l->wg_denoms[0]) == 0 && (n % mmp->l->wg_denoms[1]) == 0) || (!ctx->device->mul_mat_m && !ctx->device->mul_mat_s)) { + return aligned ? mmp->a_l : mmp->l; + } + if ((ctx->device->mul_mat_m && (m % mmp->m->wg_denoms[0]) == 0 && (n % mmp->m->wg_denoms[1]) == 0) || !ctx->device->mul_mat_s) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_s : mmp->s; + } + + if ((ctx->device->mul_mat_s && (m <= 32 || n <= 32)) || (!ctx->device->mul_mat_m && !ctx->device->mul_mat_l)) { + return aligned ? mmp->a_s : mmp->s; + } + if ((ctx->device->mul_mat_m && (m <= 64 || n <= 64)) || !ctx->device->mul_mat_l) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_l : mmp->l; +} + +static uint32_t ggml_vk_guess_matmul_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ")"); + return ggml_vk_guess_matmul_pipeline(ctx, mmp, m, n, true)->align; +} + +static void ggml_vk_matmul( + ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, + vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& split_k_buffer, + uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, + uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, + uint32_t split_k, uint32_t batch, uint32_t ne02, uint32_t ne12, uint32_t broadcast2, uint32_t broadcast3) { + VK_LOG_DEBUG("ggml_vk_matmul(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), split_k: (" << (split_k_buffer.buffer != nullptr ? split_k_buffer.buffer->buffer : VK_NULL_HANDLE) << ", " << split_k_buffer.offset << ", " << split_k_buffer.size << "), m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", split_k: " << split_k << ", batch: " << batch << ", ne02: " << ne02 << ", ne12: " << ne12 << ", broadcast2: " << broadcast2 << ", broadcast3: " << broadcast3 << ")"); + ggml_vk_sync_buffers(subctx); + if (split_k == 1) { + const vk_mat_mat_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, k, ne02, ne12, broadcast2, broadcast3 }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d }, sizeof(vk_mat_mat_push_constants), &pc, { m, n, batch }); + return; + } + + GGML_ASSERT(batch_stride_d == m * n); + + const vk_mat_mat_push_constants pc1 = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, CEIL_DIV(k, split_k), ne02, ne12, broadcast2, broadcast3 }; + // Make sure enough workgroups get assigned for split k to work + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, split_k_buffer }, sizeof(vk_mat_mat_push_constants), &pc1, { (CEIL_DIV(m, pipeline->wg_denoms[0]) * pipeline->wg_denoms[0]) * split_k, n, batch }); + ggml_vk_sync_buffers(subctx); + const std::array pc2 = { (uint32_t)(m * n * batch), split_k }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_matmul_split_k_reduce, { split_k_buffer, d }, pc2.size() * sizeof(uint32_t), pc2.data(), { m * n * batch, 1, 1 }); +} + +static vk_pipeline ggml_vk_guess_matmul_id_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, bool aligned) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline(" << m << ", " << n << ", " << aligned << ")"); + + if (ctx->device->coopmat2) { + if ((ctx->device->mul_mat_id_l && (m % mmp->l->wg_denoms[0]) == 0 && (n % mmp->l->wg_denoms[1]) == 0) || (!ctx->device->mul_mat_id_m && !ctx->device->mul_mat_id_s)) { + return aligned ? mmp->a_l : mmp->l; + } + if ((ctx->device->mul_mat_id_m && (m % mmp->m->wg_denoms[0]) == 0 && (n % mmp->m->wg_denoms[1]) == 0) || !ctx->device->mul_mat_id_s) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_s : mmp->s; + } + + if ((ctx->device->mul_mat_id_s && (m <= 32 || n <= 32)) || (!ctx->device->mul_mat_id_m && !ctx->device->mul_mat_id_l)) { + return aligned ? mmp->a_s : mmp->s; + } + if ((ctx->device->mul_mat_id_m && (m <= 64 || n <= 64)) || !ctx->device->mul_mat_id_l) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_l : mmp->l; +} + +static uint32_t ggml_vk_guess_matmul_id_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ")"); + return ggml_vk_guess_matmul_id_pipeline(ctx, mmp, m, n, true)->align; +} + +static void ggml_vk_matmul_id( + ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, + vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& ids, + uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, + uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, + uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11) { + VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), " << + "m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " << + "batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " << + "n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")"); + ggml_vk_sync_buffers(subctx); + const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, + nei0, nei1, nbi1, ne11 }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids }, sizeof(vk_mat_mat_id_push_constants), &pc, { m, nei1, n_as }); +} + +static bool ggml_vk_dim01_contiguous(const ggml_tensor * tensor) { + return + tensor->nb[0] == ggml_type_size(tensor->type) && + tensor->nb[1] == (tensor->nb[0]*tensor->ne[0])/ggml_blck_size(tensor->type) && + tensor->nb[3] == tensor->nb[2]*tensor->ne[2]; +} + +static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src, const ggml_tensor * dst, ggml_type to) { + + // Choose "contiguous copy" shader if src/dst are contiguous + bool contig = ggml_is_contiguous(src) && (!dst || ggml_is_contiguous(dst)); + + if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F32) { + if (contig) { + return ctx->device->pipeline_contig_cpy_f32_f32; + } else { + return ctx->device->pipeline_cpy_f32_f32; + } + } + if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F16) { + if (contig) { + return ctx->device->pipeline_contig_cpy_f32_f16; + } else { + return ctx->device->pipeline_cpy_f32_f16; + } + } + if (src->type == GGML_TYPE_F16 && to == GGML_TYPE_F16) { + if (contig) { + return ctx->device->pipeline_contig_cpy_f16_f16; + } else { + return ctx->device->pipeline_cpy_f16_f16; + } + } + + std::cerr << "Missing CPY op for types: " << ggml_type_name(src->type) << " " << ggml_type_name(to) << std::endl; + GGML_ABORT("fatal error"); +} + +static void ggml_vk_cpy_to_contiguous(ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline pipeline, const ggml_tensor * tensor, vk_subbuffer&& in, vk_subbuffer&& out) { + VK_LOG_DEBUG("ggml_vk_cpy_to_contiguous((" << tensor << ", type=" << tensor->type << ", ne0=" << tensor->ne[0] << ", ne1=" << tensor->ne[1] << ", ne2=" << tensor->ne[2] << ", ne3=" << tensor->ne[3] << ", nb0=" << tensor->nb[0] << ", nb1=" << tensor->nb[1] << ", nb2=" << tensor->nb[2] << ", nb3=" << tensor->nb[3] << "), "; + std::cerr << "buffer in size=" << in.buffer->size << ", buffer out size=" << out.buffer->size << ")"); + const int tensor_type_size = ggml_type_size(tensor->type); + + const uint32_t ne = ggml_nelements(tensor); + std::array elements; + + if (ne > 262144) { + elements = { 512, 512, CEIL_DIV(ne, 262144) }; + } else if (ne > 512) { + elements = { 512, CEIL_DIV(ne, 512), 1 }; + } else { + elements = { ne, 1, 1 }; + } + + vk_op_unary_push_constants pc = { + (uint32_t)ne, + (uint32_t)tensor->ne[0], (uint32_t)tensor->ne[1], (uint32_t)tensor->ne[2], (uint32_t)tensor->ne[3], (uint32_t)tensor->nb[0] / tensor_type_size, (uint32_t)tensor->nb[1] / tensor_type_size, (uint32_t)tensor->nb[2] / tensor_type_size, (uint32_t)tensor->nb[3] / tensor_type_size, + (uint32_t)tensor->ne[0], (uint32_t)tensor->ne[1], (uint32_t)tensor->ne[2], (uint32_t)tensor->ne[3], 1 , (uint32_t)tensor->ne[0] , (uint32_t)(tensor->ne[0] * tensor->ne[1]) , (uint32_t)(tensor->ne[0] * tensor->ne[1] * tensor->ne[2]), + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }; + init_pushconst_fastdiv(pc); + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { in, out }, sizeof(vk_op_unary_push_constants), &pc, elements); +} + +static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); + GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT + GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + const uint64_t ne03 = src0->ne[3]; + + const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + const uint64_t ne13 = src1->ne[3]; + + const uint64_t ne20 = dst->ne[0]; + const uint64_t ne21 = dst->ne[1]; + + const uint64_t r2 = ne12 / ne02; + const uint64_t r3 = ne13 / ne03; + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + + vk_buffer d_Qx = nullptr; + size_t qx_buf_offset = 0; + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + + bool src0_uma = false; + bool src1_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + src0_uma = d_Qx != nullptr; + src1_uma = d_Qy != nullptr; + } + + const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); + // Reformat and convert to fp16 if src1 is non-contiguous, or for coopmat2 for better perf + const bool y_non_contig = (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || + !ggml_vk_dim01_contiguous(src1); + + const bool y_f32_kernel = src1->type == GGML_TYPE_F32 && !y_non_contig; + + vk_matmul_pipeline mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, src0->type, y_non_contig ? GGML_TYPE_F16 : src1->type, (ggml_prec)dst->op_params[0]); + + const bool qx_needs_dequant = mmp == nullptr || x_non_contig; + const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !y_f32_kernel) || y_non_contig; + + if (qx_needs_dequant) { + // Fall back to dequant + f16 mulmat + mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, GGML_TYPE_F16, y_f32_kernel ? GGML_TYPE_F32 : GGML_TYPE_F16, (ggml_prec)dst->op_params[0]); + } + + // Not implemented + GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT + + const int x_ne = ne01 * ne00; + const int y_ne = ne11 * ne10; + const int d_ne = ne11 * ne01; + + const uint32_t kpad = ggml_vk_align_size(ne10, ggml_vk_guess_matmul_pipeline_align(ctx, mmp, ne01, ne11)); + const bool aligned = ne10 == kpad && ne01 > 8 && ne11 > 8; + + vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline(ctx, mmp, ne01, ne11, aligned); + + const uint32_t split_k = ggml_vk_guess_split_k(ctx, ne01, ne11, ne10, pipeline); + + const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type); + const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); + const uint64_t x_sz = !qx_needs_dequant ? qx_sz : sizeof(ggml_fp16_t) * x_ne; + const uint64_t y_sz = y_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; + const uint64_t d_sz = sizeof(float) * d_ne; + + vk_pipeline to_fp16_vk_0 = nullptr; + vk_pipeline to_fp16_vk_1 = nullptr; + + if (x_non_contig) { + to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, GGML_TYPE_F16); + } else { + to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, src0->type); + } + if (y_non_contig) { + to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, GGML_TYPE_F16); + } else { + to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); + } + GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT + GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT + + if (dryrun) { + const uint64_t x_sz_upd = x_sz * ne02 * ne03; + const uint64_t y_sz_upd = y_sz * ne12 * ne13; + const uint64_t split_k_size = split_k > 1 ? d_sz * ne12 * ne13 * split_k : 0; + if ( + (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || + (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size) || + (split_k > 1 && split_k_size > ctx->device->max_memory_allocation_size)) { + GGML_ABORT("Requested preallocation size is too large"); + } + if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { + ctx->prealloc_size_x = x_sz_upd; + } + if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { + ctx->prealloc_size_y = y_sz_upd; + } + if (split_k > 1 && ctx->prealloc_size_split_k < split_k_size) { + ctx->prealloc_size_split_k = split_k_size; + } + + // Request descriptor sets + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + if (qx_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); + } + if (qy_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); + } + if (split_k > 1) { + ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, 1); + } + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + GGML_ASSERT(d_D->size >= d_buf_offset + d_sz * ne02 * ne03); + vk_buffer d_X; + uint64_t x_buf_offset = 0; + vk_buffer d_Y; + uint64_t y_buf_offset = 0; + if (!src0_uma) { + d_Qx = src0_buf_ctx->dev_buffer; + qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + if (!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qy != nullptr); + } + if (qx_needs_dequant) { + d_X = ctx->prealloc_x; + GGML_ASSERT(d_X->size >= x_sz * ne02 * ne03); + } else { + d_X = d_Qx; + x_buf_offset = qx_buf_offset; + GGML_ASSERT(qx_sz == x_sz); + } + if (qy_needs_dequant) { + d_Y = ctx->prealloc_y; + GGML_ASSERT(d_Y->size >= y_sz * ne02 * ne03); + } else { + d_Y = d_Qy; + y_buf_offset = qy_buf_offset; + GGML_ASSERT(qy_sz == y_sz); + } + + if (x_non_contig) { + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); + } else if (qx_needs_dequant) { + const std::vector pc = { (uint32_t)ne01, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)(ggml_nelements(src0)) }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz * ne02 * ne03 }, vk_subbuffer{ d_X, 0, x_sz * ne02 * ne03 } }, pc.size() * sizeof(uint32_t), pc.data(), { (uint32_t)(x_ne * ne02 * ne03), 1, 1}); + } + if (y_non_contig) { + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); + } + + uint32_t stride_batch_x = ne00*ne01; + uint32_t stride_batch_y = ne10*ne11; + + if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { + stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); + } + + if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { + stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); + } + + // compute + ggml_vk_matmul( + ctx, subctx, pipeline, + { d_X, x_buf_offset, x_sz * ne02 * ne03 }, { d_Y, y_buf_offset, y_sz * ne12 * ne13 }, + { d_D, d_buf_offset, d_sz * ne12 * ne13 }, { ctx->prealloc_split_k, 0, d_sz * ne12 * ne13 * split_k }, + ne01, ne11, ne10, + ne10, ne10, ne01, stride_batch_x, stride_batch_y, ne20*ne21, + split_k, ne12*ne13, ne02, ne12, r2, r3 + ); // NOLINT +} + +static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_vec_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << "),)"); + GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT + GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + const uint64_t ne03 = src0->ne[3]; + + const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + const uint64_t ne13 = src1->ne[3]; + + const uint64_t ne20 = dst->ne[0]; + const uint64_t ne21 = dst->ne[1]; + const uint64_t ne22 = dst->ne[2]; + const uint64_t ne23 = dst->ne[3]; + + const uint64_t r2 = ne12 / ne02; + const uint64_t r3 = ne13 / ne03; + + // batch_n indicates that we need to compute a few vector results, and this assumes + // ne12 and ne13 are 1. It overloads the batch_strides to hold the row strides. + GGML_ASSERT(ne11 == 1 || ne12 * ne13 == 1); + bool batch_n = ne11 > 1; + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + + vk_buffer d_Qx = nullptr; + size_t qx_buf_offset = 0; + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + + bool src0_uma = false; + bool src1_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + src0_uma = d_Qx != nullptr; + src1_uma = d_Qy != nullptr; + } + + const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); + const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); + + const bool f16_f32_kernel = src1->type == GGML_TYPE_F32; + + const bool qx_needs_dequant = x_non_contig; + const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !f16_f32_kernel) || y_non_contig; + + // Not implemented + GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT + + const uint64_t x_ne = ne01 * ne00; + const uint64_t y_ne = ne11 * ne10; + const uint64_t d_ne = ne11 * ne01; + + const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); + const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); + const uint64_t x_sz = x_non_contig ? ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment) : qx_sz; + const uint64_t y_sz = f16_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; + const uint64_t d_sz = sizeof(float) * d_ne; + + vk_pipeline to_fp16_vk_0 = nullptr; + vk_pipeline to_fp16_vk_1 = nullptr; + if (x_non_contig) { + to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, src0->type); + } + if (y_non_contig) { + to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, src1->type); + } else { + to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); + } + vk_pipeline dmmv = ggml_vk_get_dequantize_mul_mat_vec(ctx, src0->type, src1->type, ne11); + GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT + GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT + GGML_ASSERT(dmmv != nullptr); + + if (dryrun) { + const uint64_t x_sz_upd = x_sz * ne02 * ne03; + const uint64_t y_sz_upd = y_sz * ne12 * ne13; + if ( + (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || + (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { + GGML_ABORT("Requested preallocation size is too large"); + } + if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { + ctx->prealloc_size_x = x_sz_upd; + } + if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { + ctx->prealloc_size_y = y_sz_upd; + } + + // Request descriptor sets + if (qx_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); + } + if (qy_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); + } + ggml_pipeline_request_descriptor_sets(ctx->device, dmmv, 1); + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + vk_buffer d_X; + uint64_t x_buf_offset = 0; + vk_buffer d_Y; + uint64_t y_buf_offset = 0; + if(!src0_uma) { + d_Qx = src0_buf_ctx->dev_buffer; + qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + if(!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qy != nullptr); + } + if (qx_needs_dequant) { + d_X = ctx->prealloc_x; + } else { + d_X = d_Qx; + x_buf_offset = qx_buf_offset; + GGML_ASSERT(qx_sz == x_sz); + } + if (qy_needs_dequant) { + d_Y = ctx->prealloc_y; + } else { + d_Y = d_Qy; + y_buf_offset = qy_buf_offset; + GGML_ASSERT(qy_sz == y_sz); + } + + if (x_non_contig) { + GGML_ASSERT(x_sz == ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment)); + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); + } + if (y_non_contig) { + GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); + } + + // For batch_n, the A matrix is the same for each batch, and B/D use the row stride as the batch stride + uint32_t stride_batch_x = batch_n ? 0 : ne00*ne01; + uint32_t stride_batch_y = batch_n ? ne10 : (ne10*ne11); + uint32_t stride_batch_d = batch_n ? ne20 : (ne20*ne21); + + if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { + stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); + } + + if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { + stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); + } + + const uint32_t max_groups_x = ctx->device->properties.limits.maxComputeWorkGroupCount[0]; + + uint32_t groups_x = ne01; + uint32_t groups_z = 1; + + if (ne01 > max_groups_x) { + groups_z = 64; + groups_x = CEIL_DIV(groups_x, groups_z); + } + + // compute + const vk_mat_vec_push_constants pc = { + (uint32_t)ne00, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne01, + stride_batch_x, stride_batch_y, stride_batch_d, + (uint32_t)ne02, (uint32_t)ne12, (uint32_t)r2, (uint32_t)r3, + }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, dmmv, + { vk_subbuffer{ d_X, x_buf_offset, x_sz * ne02 * ne03 }, vk_subbuffer{ d_Y, y_buf_offset, y_sz * ne12 * ne13 }, vk_subbuffer{ d_D, d_buf_offset, d_sz * ne22 * ne23} }, + sizeof(vk_mat_vec_push_constants), &pc, { groups_x, (uint32_t)(ne12 * ne13), groups_z }); +} + +static void ggml_vk_mul_mat_vec_p021_f16_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_p021_f16_f32(" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); + GGML_ASSERT(ggml_is_permuted(src0) && ggml_is_permuted(src1)); + GGML_ASSERT(src0->nb[0] <= src0->nb[1] && src0->nb[2] <= src0->nb[3]); // NOLINT + GGML_ASSERT(src1->nb[0] <= src1->nb[1] && src1->nb[2] <= src1->nb[3]); // NOLINT + GGML_ASSERT(src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + // const uint64_t ne03 = src0->ne[3]; + + const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + // const uint64_t ne13 = src1->ne[3]; + + GGML_ASSERT(ne11 == 1); + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + + bool src1_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + src1_uma = d_Qy != nullptr; + } + + const uint64_t x_ne = ne00 * ne01 * ne02; + const uint64_t y_ne = ne10 * ne11 * ne12; + const uint64_t d_ne = ne01 * ne11 * ne12; + + const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); + const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); + const uint64_t d_sz = sizeof(float) * d_ne; + + if (dryrun) { + // Request descriptor sets + ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_mul_mat_vec_p021_f16_f32, 1); + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + vk_buffer d_Qx = src0_buf_ctx->dev_buffer; + const uint64_t qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + if (!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + + const uint64_t qy_buffer_offset = (qy_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; + const uint64_t qy_shader_offset = qy_buf_offset - qy_buffer_offset; + + const uint64_t d_buffer_offset = (d_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; + const uint64_t d_shader_offset = d_buf_offset - d_buffer_offset; + + // compute + const std::array pc = { (uint32_t)ne00, (uint32_t)ne01, (uint32_t)ne02, (uint32_t)ne12, (uint32_t)(qy_shader_offset / ggml_type_size(src1->type)), (uint32_t)(d_shader_offset / ggml_type_size(dst->type)) }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_mul_mat_vec_p021_f16_f32, { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz }, vk_subbuffer{ d_Qy, qy_buffer_offset, qy_sz + qy_shader_offset }, vk_subbuffer{ d_D, d_buffer_offset, d_sz + d_shader_offset } }, 6 * sizeof(uint32_t), &pc, { 1, (uint32_t)ne01, (uint32_t)ne12 }); +} + +static void ggml_vk_mul_mat_vec_nc_f16_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_nc_f16_f32((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); + GGML_ASSERT(!ggml_is_transposed(src0)); + GGML_ASSERT(!ggml_is_transposed(src1)); + GGML_ASSERT(!ggml_is_permuted(src0)); + GGML_ASSERT(src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + // const uint64_t ne03 = src0->ne[3]; + + const uint64_t nb01 = src0->nb[1]; + const uint64_t nb02 = src0->nb[2]; + + // const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + // const uint64_t ne13 = src1->ne[3]; + + GGML_ASSERT(ne11 == 1); + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + + bool src1_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + src1_uma = d_Qy != nullptr; + } + + const uint64_t d_ne = ne01 * ne11 * ne12; + + const uint32_t row_stride_x = nb01 / sizeof(ggml_fp16_t); + const uint32_t channel_stride_x = nb02 / sizeof(ggml_fp16_t); + + const uint64_t qx_sz = ggml_nbytes(src0); + const uint64_t qy_sz = ggml_nbytes(src1); + const uint64_t d_sz = sizeof(float) * d_ne; + + if (dryrun) { + // Request descriptor sets + ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_mul_mat_vec_nc_f16_f32, 1); + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + vk_buffer d_Qx = src0_buf_ctx->dev_buffer; + const uint64_t qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + if (!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + + const uint64_t qy_buffer_offset = (qy_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; + const uint64_t qy_shader_offset = qy_buf_offset - qy_buffer_offset; + + const uint64_t d_buffer_offset = (d_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; + const uint64_t d_shader_offset = d_buf_offset - d_buffer_offset; + + // compute + const std::array pc = { (uint32_t)ne00, (uint32_t)ne01, row_stride_x, channel_stride_x, (uint32_t)(ne12 / ne02), (uint32_t)(qy_shader_offset / ggml_type_size(src1->type)), (uint32_t)(d_shader_offset / ggml_type_size(dst->type)) }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_mul_mat_vec_nc_f16_f32, + { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz }, vk_subbuffer{ d_Qy, qy_buffer_offset, qy_sz + qy_shader_offset }, vk_subbuffer{ d_D, d_buffer_offset, d_sz + d_shader_offset } }, 7 * sizeof(uint32_t), &pc, { 1, (uint32_t)ne01, (uint32_t)ne12 }); +} + +static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat(" << src0 << ", " << src1 << ", " << dst << ")"); + if (src0->type == GGML_TYPE_F16 && ggml_is_permuted(src0) && ggml_is_permuted(src1) && dst->ne[1] == 1 && + // detect 0213 permutation, and batch size of 1 + src0->nb[0] <= src0->nb[2] && + src0->nb[2] <= src0->nb[1] && + src0->nb[1] <= src0->nb[3] && + src1->nb[0] <= src1->nb[2] && + src1->nb[2] <= src1->nb[1] && + src1->nb[1] <= src1->nb[3] && + src0->ne[3] == 1 && + src1->ne[3] == 1) { + ggml_vk_mul_mat_vec_p021_f16_f32(ctx, subctx, src0, src1, dst, dryrun); + } else if (src0->type == GGML_TYPE_F16 && !ggml_is_contiguous(src0) && !ggml_is_transposed(src1) && dst->ne[1] == 1 && + !ggml_is_permuted(src0) && !ggml_is_permuted(src1)) { + ggml_vk_mul_mat_vec_nc_f16_f32(ctx, subctx, src0, src1, dst, dryrun); + // mul_mat_vec supports batching ne12*ne13 when ne11==1, or treating ne11 as the batch size (up to four) + // when ne12 and ne13 are one. + } else if ((dst->ne[1] == 1 || (dst->ne[1] <= mul_mat_vec_max_cols && src1->ne[2] * src1->ne[3] == 1)) && + (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type))) { + ggml_vk_mul_mat_vec_q_f16(ctx, subctx, src0, src1, dst, dryrun); + } else { + ggml_vk_mul_mat_q_f16(ctx, subctx, src0, src1, dst, dryrun); + } +} + +static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_id_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << ids << ", name=" << ids->name << ", type=" << ids->type << ", ne0=" << ids->ne[0] << ", ne1=" << ids->ne[1] << ", ne2=" << ids->ne[2] << ", ne3=" << ids->ne[3] << ", nb0=" << ids->nb[0] << ", nb1=" << ids->nb[1] << ", nb2=" << ids->nb[2] << ", nb3=" << ids->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3] << "),)"); + GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT + GGML_ASSERT(ids->type == GGML_TYPE_I32); + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + const uint64_t ne03 = src0->ne[3]; + + const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + const uint64_t ne13 = src1->ne[3]; + + const uint64_t nei0 = ids->ne[0]; + const uint64_t nei1 = ids->ne[1]; + GGML_ASSERT(nei0 * nei1 <= 3072); + + const uint32_t nbi1 = ids->nb[1]; + const uint32_t nbi2 = ids->nb[2]; + + const uint64_t ne20 = dst->ne[0]; + const uint64_t ne21 = dst->ne[1]; + const uint64_t ne22 = dst->ne[2]; + const uint64_t ne23 = dst->ne[3]; + + const uint64_t n_as = ne02; + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + ggml_backend_vk_buffer_context * ids_buf_ctx = (ggml_backend_vk_buffer_context *)ids->buffer->context; + + vk_buffer d_Qx = nullptr; + size_t qx_buf_offset = 0; + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + vk_buffer d_ids = nullptr; + size_t ids_buf_offset = 0; + + bool src0_uma = false; + bool src1_uma = false; + bool ids_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + ggml_vk_host_get(ctx->device, ids->data, d_ids, ids_buf_offset); + src0_uma = d_Qx != nullptr; + src1_uma = d_Qy != nullptr; + ids_uma = d_ids != nullptr; + } + + const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); + const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); + + const bool y_f32_kernel = src1->type == GGML_TYPE_F32 && !y_non_contig; + + vk_matmul_pipeline mmp = ggml_vk_get_mul_mat_mat_id_pipeline(ctx, src0->type, y_non_contig ? GGML_TYPE_F16 : src1->type, (ggml_prec)dst->op_params[0]); + + const bool qx_needs_dequant = mmp == nullptr || x_non_contig; + const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !y_f32_kernel) || y_non_contig; + + if (qx_needs_dequant) { + GGML_ABORT("fatal error"); + } + + // Not implemented + GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT + + const uint64_t x_ne = ne01 * ne00; + const uint64_t y_ne = ne11 * ne10; + const uint64_t d_ne = ne21 * ne20; + + const uint32_t kpad = ggml_vk_align_size(ne10, ggml_vk_guess_matmul_id_pipeline_align(ctx, mmp, ne01, nei1)); + const bool aligned = ne10 == kpad && ne01 > 8 && nei1 > 8; + + vk_pipeline pipeline = ggml_vk_guess_matmul_id_pipeline(ctx, mmp, ne01, nei1, aligned); + + const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type); + const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); + const uint64_t x_sz = !qx_needs_dequant ? qx_sz : sizeof(ggml_fp16_t) * x_ne; + const uint64_t y_sz = y_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; + const uint64_t ids_sz = nbi2; + const uint64_t d_sz = sizeof(float) * d_ne; + + vk_pipeline to_fp16_vk_0 = nullptr; + vk_pipeline to_fp16_vk_1 = nullptr; + + if (x_non_contig) { + to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, GGML_TYPE_F16); + } else { + to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, src0->type); + } + if (y_non_contig) { + to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, GGML_TYPE_F16); + } else { + to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); + } + GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT + GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT + + if (dryrun) { + const uint64_t x_sz_upd = x_sz * ne02 * ne03; + const uint64_t y_sz_upd = y_sz * ne12 * ne13; + if ( + (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || + (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { + GGML_ABORT("Requested preallocation size is too large"); + } + if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { + ctx->prealloc_size_x = x_sz_upd; + } + if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { + ctx->prealloc_size_y = y_sz_upd; + } + + // Request descriptor sets + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + if (qx_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); + } + if (qy_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); + } + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + vk_buffer d_X; + uint64_t x_buf_offset = 0; + vk_buffer d_Y; + uint64_t y_buf_offset = 0; + if (!src0_uma) { + d_Qx = src0_buf_ctx->dev_buffer; + qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + if (!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qy != nullptr); + } + if (!ids_uma) { + d_ids = ids_buf_ctx->dev_buffer; + ids_buf_offset = vk_tensor_offset(ids) + ids->view_offs; + GGML_ASSERT(d_ids != nullptr); + } + if (qx_needs_dequant) { + d_X = ctx->prealloc_x; + GGML_ASSERT(d_X->size >= x_sz * ne02 * ne03); + } else { + d_X = d_Qx; + x_buf_offset = qx_buf_offset; + GGML_ASSERT(qx_sz == x_sz); + } + if (qy_needs_dequant) { + d_Y = ctx->prealloc_y; + GGML_ASSERT(d_Y->size >= y_sz * ne02 * ne03); + } else { + d_Y = d_Qy; + y_buf_offset = qy_buf_offset; + GGML_ASSERT(qy_sz == y_sz); + } + + if (x_non_contig) { + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); + } else if (qx_needs_dequant) { + const std::vector pc = { (uint32_t)ne01, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)(ggml_nelements(src0)) }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, + { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz * ne02 * ne03 }, vk_subbuffer{ d_X, 0, x_sz * ne02 * ne03 } }, pc.size() * sizeof(uint32_t), pc.data(), { (uint32_t)(x_ne * ne02 * ne03), 1, 1}); + } + if (y_non_contig) { + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); + } + + uint32_t stride_batch_x = ne00*ne01; + uint32_t stride_batch_y = ne10*ne11; + + if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { + stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); + } + + if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { + stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); + } + + // compute + ggml_vk_matmul_id( + ctx, subctx, pipeline, + { d_X, x_buf_offset, x_sz * ne02 * ne03 }, { d_Y, y_buf_offset, y_sz * ne12 * ne13 }, + { d_D, d_buf_offset, d_sz * ne22 * ne23 }, { d_ids, ids_buf_offset, ids_sz }, + ne01, ne21, ne10, ne10, ne10, ne01, + stride_batch_x, stride_batch_y, ne20*ne21, + n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11 + ); // NOLINT +} + +static void ggml_vk_mul_mat_vec_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_vec_id_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + std::cerr << "), (" << ids << ", name=" << ids->name << ", type=" << ids->type << ", ne0=" << ids->ne[0] << ", ne1=" << ids->ne[1] << ", ne2=" << ids->ne[2] << ", ne3=" << ids->ne[3] << ", nb0=" << ids->nb[0] << ", nb1=" << ids->nb[1] << ", nb2=" << ids->nb[2] << ", nb3=" << ids->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); + GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT + GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT + GGML_ASSERT(ids->type == GGML_TYPE_I32); + + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + const uint64_t ne03 = src0->ne[3]; + + const uint64_t ne10 = src1->ne[0]; + const uint64_t ne11 = src1->ne[1]; + const uint64_t ne12 = src1->ne[2]; + const uint64_t ne13 = src1->ne[3]; + + const uint64_t nei0 = ids->ne[0]; + const uint64_t nei1 = ids->ne[1]; + + const uint64_t nbi2 = ids->nb[2]; + + GGML_ASSERT(nei1 == 1); + + const uint64_t ne20 = dst->ne[0]; + const uint64_t ne21 = dst->ne[1]; + const uint64_t ne22 = dst->ne[2]; + const uint64_t ne23 = dst->ne[3]; + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + ggml_backend_vk_buffer_context * ids_buf_ctx = (ggml_backend_vk_buffer_context *)ids->buffer->context; + + vk_buffer d_Qx = nullptr; + size_t qx_buf_offset = 0; + vk_buffer d_Qy = nullptr; + size_t qy_buf_offset = 0; + vk_buffer d_ids = nullptr; + size_t ids_buf_offset = 0; + + bool src0_uma = false; + bool src1_uma = false; + bool ids_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); + ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); + ggml_vk_host_get(ctx->device, ids->data, d_ids, ids_buf_offset); + src0_uma = d_Qx != nullptr; + src1_uma = d_Qy != nullptr; + ids_uma = d_ids != nullptr; + } + + const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); + const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); + + const bool f16_f32_kernel = src1->type == GGML_TYPE_F32; + + const bool qx_needs_dequant = x_non_contig; + const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !f16_f32_kernel) || y_non_contig; + + // Not implemented + GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT + + const uint64_t x_ne = ne01 * ne00; + const uint64_t y_ne = ne11 * ne10; + const uint64_t d_ne = ne21 * ne20; + + const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); + const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); + const uint64_t x_sz = x_non_contig ? ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment) : qx_sz; + const uint64_t y_sz = f16_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; + const uint64_t ids_sz = nbi2; + const uint64_t d_sz = sizeof(float) * d_ne; + + vk_pipeline to_fp16_vk_0 = nullptr; + vk_pipeline to_fp16_vk_1 = nullptr; + if (x_non_contig) { + to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, src0->type); + } + if (y_non_contig) { + to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, src1->type); + } else { + to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); + } + vk_pipeline dmmv = ggml_vk_get_dequantize_mul_mat_vec_id(ctx, src0->type, src1->type); + GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT + GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT + GGML_ASSERT(dmmv != nullptr); + + if (dryrun) { + const uint64_t x_sz_upd = x_sz * ne02 * ne03; + const uint64_t y_sz_upd = y_sz * ne12 * ne13; + if ( + (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || + (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { + GGML_ABORT("Requested preallocation size is too large"); + } + if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { + ctx->prealloc_size_x = x_sz_upd; + } + if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { + ctx->prealloc_size_y = y_sz_upd; + } + + // Request descriptor sets + if (qx_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); + } + if (qy_needs_dequant) { + ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); + } + ggml_pipeline_request_descriptor_sets(ctx->device, dmmv, 1); + return; + } + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + GGML_ASSERT(d_D != nullptr); + vk_buffer d_X; + uint64_t x_buf_offset = 0; + vk_buffer d_Y; + uint64_t y_buf_offset = 0; + if(!src0_uma) { + d_Qx = src0_buf_ctx->dev_buffer; + qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_Qx != nullptr); + } + if(!src1_uma) { + d_Qy = src1_buf_ctx->dev_buffer; + qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Qy != nullptr); + } + if(!ids_uma) { + d_ids = ids_buf_ctx->dev_buffer; + ids_buf_offset = vk_tensor_offset(ids) + ids->view_offs; + GGML_ASSERT(d_ids != nullptr); + } + if (qx_needs_dequant) { + d_X = ctx->prealloc_x; + } else { + d_X = d_Qx; + x_buf_offset = qx_buf_offset; + GGML_ASSERT(qx_sz == x_sz); + } + if (qy_needs_dequant) { + d_Y = ctx->prealloc_y; + } else { + d_Y = d_Qy; + y_buf_offset = qy_buf_offset; + GGML_ASSERT(qy_sz == y_sz); + } + + if (x_non_contig) { + GGML_ASSERT(x_sz == ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment)); + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); + } + if (y_non_contig) { + GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); + ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); + } + + uint32_t stride_batch_y = ne10*ne11; + + if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { + stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); + } + + const uint32_t max_groups_x = ctx->device->properties.limits.maxComputeWorkGroupCount[0]; + + uint32_t groups_x = ne01; + uint32_t groups_z = 1; + + if (ne01 > max_groups_x) { + groups_z = 64; + groups_x = CEIL_DIV(groups_x, groups_z); + } + + // compute + const vk_mat_vec_id_push_constants pc = { + (uint32_t)ne00, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne01, + (uint32_t)x_ne, stride_batch_y, (uint32_t)(ne20*ne21), + (uint32_t)nei0, (uint32_t)ne11, + }; + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, dmmv, + { vk_subbuffer{ d_X, x_buf_offset, x_sz * ne02 * ne03 }, + vk_subbuffer{ d_Y, y_buf_offset, y_sz * ne12 * ne13 }, vk_subbuffer{ d_D, d_buf_offset, d_sz * ne22 * ne23}, vk_subbuffer{ d_ids, ids_buf_offset, ids_sz } }, + sizeof(vk_mat_vec_id_push_constants), &pc, { groups_x, (uint32_t)nei0, groups_z }); +} + +static void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_mul_mat_id(" << src0 << ", " << src1 << ", " << src2 << ", " << dst << ")"); + if (src2->ne[1] == 1 && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type))) { + ggml_vk_mul_mat_vec_id_q_f16(ctx, subctx, src0, src1, src2, dst, dryrun); + } else { + ggml_vk_mul_mat_id_q_f16(ctx, subctx, src0, src1, src2, dst, dryrun); + } +} + +static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, ggml_tensor * dst, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_flash_attn((" << q << ", name=" << q->name << ", type=" << q->type << ", ne0=" << q->ne[0] << ", ne1=" << q->ne[1] << ", ne2=" << q->ne[2] << ", ne3=" << q->ne[3] << ", nb0=" << q->nb[0] << ", nb1=" << q->nb[1] << ", nb2=" << q->nb[2] << ", nb3=" << q->nb[3]; + std::cerr << "), (" << k << ", name=" << k->name << ", type=" << k->type << ", ne0=" << k->ne[0] << ", ne1=" << k->ne[1] << ", ne2=" << k->ne[2] << ", ne3=" << k->ne[3] << ", nb0=" << k->nb[0] << ", nb1=" << k->nb[1] << ", nb2=" << k->nb[2] << ", nb3=" << k->nb[3]; + std::cerr << "), (" << v << ", name=" << v->name << ", type=" << v->type << ", ne0=" << v->ne[0] << ", ne1=" << v->ne[1] << ", ne2=" << v->ne[2] << ", ne3=" << v->ne[3] << ", nb0=" << v->nb[0] << ", nb1=" << v->nb[1] << ", nb2=" << v->nb[2] << ", nb3=" << v->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne) + GGML_TENSOR_LOCALS(size_t, nbq, q, nb) + GGML_TENSOR_LOCALS(int64_t, nek, k, ne) + GGML_TENSOR_LOCALS(size_t, nbk, k, nb) + GGML_TENSOR_LOCALS(int64_t, nev, v, ne) + GGML_TENSOR_LOCALS(size_t, nbv, v, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + + const uint32_t nem1 = mask ? mask->ne[1] : 0; + const uint32_t nbm1 = mask ? mask->nb[1] : 0; + + const uint32_t D = neq0; + const uint32_t N = neq1; + const uint32_t KV = nek1; + + GGML_ASSERT(ne0 == D); + GGML_ASSERT(ne2 == N); + + // input tensor rows must be contiguous + GGML_ASSERT(nbq0 == ggml_type_size(q->type)); + GGML_ASSERT(nbk0 == ggml_type_size(k->type)); + GGML_ASSERT(nbv0 == ggml_type_size(v->type)); + + GGML_ASSERT(neq0 == D); + GGML_ASSERT(nek0 == D); + GGML_ASSERT(nev0 == D); + + GGML_ASSERT(neq1 == N); + GGML_ASSERT(nev0 == D); + + GGML_ASSERT(nev1 == nek1); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + assert(dst->type == GGML_TYPE_F32); + assert(q->type == GGML_TYPE_F32); + assert(k->type == v->type); + + vk_pipeline *pipelines; + // XXX TODO other backends may be changing accumulator precision to default to f32 soon + bool f32acc = dst->op_params[3] == GGML_PREC_F32; + bool small_rows = N <= flash_attention_num_small_rows; + switch (D) { + case 64: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D64[k->type][f32acc][small_rows][0]; break; + case 80: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D80[k->type][f32acc][small_rows][0]; break; + case 96: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D96[k->type][f32acc][small_rows][0]; break; + case 112: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D112[k->type][f32acc][small_rows][0]; break; + case 128: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D128[k->type][f32acc][small_rows][0]; break; + case 256: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D256[k->type][f32acc][small_rows][0]; break; + default: + assert(!"unsupported D value"); + return; + } + assert(pipelines); + + bool aligned = (KV % pipelines[1]->align) == 0; + vk_pipeline pipeline = pipelines[aligned]; + assert(pipeline); + + if (dryrun) { + // Request descriptor sets + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + return; + } + + float scale = 1.0f; + float max_bias = 0.0f; + float logit_softcap = 0.0f; + + memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + + if (logit_softcap != 0) { + scale /= logit_softcap; + } + + const uint32_t n_head_kv = neq2; + const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head_kv)); + const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); + const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + + ggml_vk_sync_buffers(subctx); + + vk_buffer d_Q = nullptr, d_K = nullptr, d_V = nullptr, d_D = nullptr, d_M = nullptr; + size_t q_buf_offset = 0, k_buf_offset = 0, v_buf_offset = 0, d_buf_offset = 0, m_buf_offset = 0; + + bool Q_uma = false, K_uma = false, V_uma = false, D_uma = false, M_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, q->data, d_Q, q_buf_offset); + ggml_vk_host_get(ctx->device, k->data, d_K, q_buf_offset); + ggml_vk_host_get(ctx->device, v->data, d_V, q_buf_offset); + ggml_vk_host_get(ctx->device, dst->data, d_D, q_buf_offset); + Q_uma = d_Q != nullptr; + K_uma = d_K != nullptr; + V_uma = d_V != nullptr; + D_uma = d_D != nullptr; + if (mask) { + ggml_vk_host_get(ctx->device, mask->data, d_M, q_buf_offset); + M_uma = d_M != nullptr; + } + } + + + ggml_backend_vk_buffer_context * d_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * q_buf_ctx = (ggml_backend_vk_buffer_context *)q->buffer->context; + ggml_backend_vk_buffer_context * k_buf_ctx = (ggml_backend_vk_buffer_context *)k->buffer->context; + ggml_backend_vk_buffer_context * v_buf_ctx = (ggml_backend_vk_buffer_context *)v->buffer->context; + + if (!Q_uma) { + d_Q = q_buf_ctx->dev_buffer; + q_buf_offset = vk_tensor_offset(q) + q->view_offs; + } + if (!K_uma) { + d_K = k_buf_ctx->dev_buffer; + k_buf_offset = vk_tensor_offset(k) + k->view_offs; + } + if (!V_uma) { + d_V = v_buf_ctx->dev_buffer; + v_buf_offset = vk_tensor_offset(v) + v->view_offs; + } + if (!D_uma) { + d_D = d_buf_ctx->dev_buffer; + d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + } + + if (!M_uma) { + d_M = d_Q; + m_buf_offset = q_buf_offset; + if (mask) { + ggml_backend_vk_buffer_context * m_buf_ctx = (ggml_backend_vk_buffer_context*)mask->buffer->context; + d_M = m_buf_ctx->dev_buffer; + m_buf_offset = vk_tensor_offset(mask) + mask->view_offs; + } + } + + const vk_flash_attn_push_constants pc = { N, KV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3, (uint32_t)neq2, (uint32_t)neq3, (uint32_t)nek2, (uint32_t)nek3, (uint32_t)nev2, (uint32_t)nev3, nem1, (uint32_t)nbq2, (uint32_t)nbq3, (uint32_t)nbk2, (uint32_t)nbk3, (uint32_t)nbv2, (uint32_t)nbv3, nbm1, scale, max_bias, logit_softcap, mask != nullptr, n_head_log2, m0, m1 }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + vk_subbuffer{d_Q, q_buf_offset, VK_WHOLE_SIZE}, + vk_subbuffer{d_K, k_buf_offset, VK_WHOLE_SIZE}, + vk_subbuffer{d_V, v_buf_offset, VK_WHOLE_SIZE}, + vk_subbuffer{d_M, m_buf_offset, VK_WHOLE_SIZE}, + vk_subbuffer{d_D, d_buf_offset, VK_WHOLE_SIZE}, + }, + sizeof(vk_flash_attn_push_constants), &pc, { (uint32_t)neq1, (uint32_t)neq2, (uint32_t)neq3 }); +} + +static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, ggml_op op) { + switch (op) { + case GGML_OP_GET_ROWS: + GGML_ASSERT(src1->type == GGML_TYPE_I32); + if (dst->type == GGML_TYPE_F16) { + return ctx->device->pipeline_get_rows[src0->type]; + } + if (dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_get_rows_f32[src0->type]; + } + return nullptr; + case GGML_OP_ACC: + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_acc_f32; + } + return nullptr; + case GGML_OP_ADD: + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_add_f32_norepeat : ctx->device->pipeline_add_f32; + } + if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) { + return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_add_f16_f32_f16_norepeat : ctx->device->pipeline_add_f16_f32_f16; + } + return nullptr; + case GGML_OP_MUL: + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_mul_f32_norepeat : ctx->device->pipeline_mul_f32; + } + return nullptr; + case GGML_OP_DIV: + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_div_f32_norepeat : ctx->device->pipeline_div_f32; + } + return nullptr; + case GGML_OP_CONCAT: + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_concat_f32; + } + if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { + return ctx->device->pipeline_concat_f16; + } + if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_I32 && dst->type == GGML_TYPE_I32) { + return ctx->device->pipeline_concat_i32; + } + return nullptr; + case GGML_OP_UPSCALE: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_upscale_f32; + } + return nullptr; + case GGML_OP_SCALE: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_scale_f32; + } + return nullptr; + case GGML_OP_SQR: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_sqr_f32; + } + return nullptr; + case GGML_OP_SIN: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_sin_f32; + } + return nullptr; + case GGML_OP_COS: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_cos_f32; + } + return nullptr; + case GGML_OP_CLAMP: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_clamp_f32; + } + return nullptr; + case GGML_OP_PAD: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_pad_f32; + } + return nullptr; + case GGML_OP_REPEAT: + if (ggml_type_size(src0->type) == sizeof(float) && ggml_type_size(dst->type) == sizeof(float)) { + return ctx->device->pipeline_repeat_f32; + } + return nullptr; + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_DUP: + return ggml_vk_get_cpy_pipeline(ctx, src0, dst, dst->type); + case GGML_OP_NORM: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_norm_f32; + } + return nullptr; + case GGML_OP_GROUP_NORM: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_group_norm_f32; + } + return nullptr; + case GGML_OP_RMS_NORM: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_rms_norm_f32; + } + return nullptr; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(dst)) { + case GGML_UNARY_OP_SILU: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_silu_f32; + } + break; + case GGML_UNARY_OP_GELU: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_gelu_f32; + } + break; + case GGML_UNARY_OP_GELU_QUICK: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_gelu_quick_f32; + } + break; + case GGML_UNARY_OP_RELU: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_relu_f32; + } + break; + case GGML_UNARY_OP_TANH: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_tanh_f32; + } + break; + default: + break; + } + return nullptr; + case GGML_OP_DIAG_MASK_INF: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_diag_mask_inf_f32; + } + return nullptr; + case GGML_OP_SOFT_MAX: + GGML_ASSERT(!src1 || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); + + if (src0->type == GGML_TYPE_F32 && (src1 == nullptr || src1->type == GGML_TYPE_F32) && dst->type == GGML_TYPE_F32) { + return src0->ne[0] > 1024 ? ctx->device->pipeline_soft_max_f32_wg512 : ctx->device->pipeline_soft_max_f32; + } + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) { + return src0->ne[0] > 1024 ? ctx->device->pipeline_soft_max_f32_f16_wg512 : ctx->device->pipeline_soft_max_f32_f16; + } + return nullptr; + case GGML_OP_ROPE: + { + const int mode = ((const int32_t *) dst->op_params)[2]; + const bool is_neox = mode & GGML_ROPE_TYPE_NEOX; + + if (is_neox) { + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_rope_neox_f32; + } + if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { + return ctx->device->pipeline_rope_neox_f16; + } + } else { + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_rope_norm_f32; + } + if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { + return ctx->device->pipeline_rope_norm_f16; + } + } + return nullptr; + } + case GGML_OP_ARGSORT: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_I32) { + return ctx->device->pipeline_argsort_f32; + } + return nullptr; + case GGML_OP_SUM_ROWS: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_sum_rows_f32; + } + return nullptr; + case GGML_OP_IM2COL: + if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_im2col_f32; + } + if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) { + return ctx->device->pipeline_im2col_f32_f16; + } + return nullptr; + case GGML_OP_TIMESTEP_EMBEDDING: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_timestep_embedding_f32; + } + return nullptr; + case GGML_OP_POOL_2D: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_pool2d_f32; + } + return nullptr; + case GGML_OP_RWKV_WKV6: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_rwkv_wkv6_f32; + } + return nullptr; + case GGML_OP_LEAKY_RELU: + if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return ctx->device->pipeline_leaky_relu_f32; + } + return nullptr; + default: + return nullptr; + } + + GGML_UNUSED(src2); +} + +static bool ggml_vk_op_supports_incontiguous(ggml_op op) { + switch (op) { + case GGML_OP_CPY: + case GGML_OP_GET_ROWS: + case GGML_OP_ADD: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_REPEAT: + return true; + default: + return false; + } +} + +static uint32_t get_misalign_bytes(ggml_backend_vk_context * ctx, const ggml_tensor * t) +{ + return ((vk_tensor_offset(t) + t->view_offs) & (ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1));; +} + +template void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, T &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { + GGML_UNUSED(p); + GGML_UNUSED(src0); + GGML_UNUSED(src1); + GGML_UNUSED(src2); + GGML_UNUSED(dst); + static_assert(!std::is_const::value, "unexpected type"); + GGML_ASSERT(!src0 || get_misalign_bytes(ctx, src0) == 0); + GGML_ASSERT(!src1 || get_misalign_bytes(ctx, src1) == 0); + GGML_ASSERT(!src2 || get_misalign_bytes(ctx, src2) == 0); + GGML_ASSERT(!dst || get_misalign_bytes(ctx, dst) == 0); +} + +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_unary_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { + const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + p.misalign_offsets = (a_offset << 16) | d_offset; + + GGML_UNUSED(src1); + GGML_UNUSED(src2); +} + +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_binary_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { + const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + const uint32_t b_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); + const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + GGML_ASSERT(dst->op != GGML_OP_GET_ROWS || (a_offset == 0 && b_offset == 0 && d_offset == 0)); + + p.misalign_offsets = (a_offset << 16) | (b_offset << 8) | d_offset; + + GGML_UNUSED(src2); +} + +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_upscale_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { + const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + p.a_offset = a_offset; + p.d_offset = d_offset; + + GGML_UNUSED(src1); + GGML_UNUSED(src2); +} + +template +static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, ggml_op op, PC&& pc, bool dryrun = false) { + VK_LOG_DEBUG("ggml_vk_op_f32((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; + if (src1 != nullptr) { + std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; + } + if (src2 != nullptr) { + std::cerr << "), (" << src2 << ", name=" << src2->name << ", type=" << src2->type << ", ne0=" << src2->ne[0] << ", ne1=" << src2->ne[1] << ", ne2=" << src2->ne[2] << ", ne3=" << src2->ne[3] << ", nb0=" << src2->nb[0] << ", nb1=" << src2->nb[1] << ", nb2=" << src2->nb[2] << ", nb3=" << src2->nb[3]; + } + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + std::cerr << "), " << ggml_op_name(op) << ", " << (dryrun ? "dryrun" : "") << ")"); + GGML_ASSERT(op == GGML_OP_GET_ROWS || (!ggml_is_quantized(src0->type) && (src1 == nullptr || !ggml_is_quantized(src1->type)))); // NOLINT + GGML_ASSERT(ggml_vk_op_supports_incontiguous(op) || ggml_vk_dim01_contiguous(src0)); // NOLINT + GGML_ASSERT(dst->buffer != nullptr); + const uint64_t ne00 = src0->ne[0]; + const uint64_t ne01 = src0->ne[1]; + const uint64_t ne02 = src0->ne[2]; + const uint64_t ne03 = src0->ne[3]; + const uint64_t ne0 = ne00 * ne01; + + const bool use_src1 = src1 != nullptr; + const uint64_t ne10 = use_src1 ? src1->ne[0] : 0; + const uint64_t ne11 = use_src1 ? src1->ne[1] : 0; + const uint64_t ne12 = use_src1 ? src1->ne[2] : 0; + const uint64_t ne13 = use_src1 ? src1->ne[3] : 0; + const uint64_t ne1 = ne10 * ne11; + // const uint64_t nb10 = use_src1 ? src1->nb[0] : 0; + + const bool use_src2 = src2 != nullptr; + const uint64_t ne20 = use_src2 ? src2->ne[0] : 0; + const uint64_t ne21 = use_src2 ? src2->ne[1] : 0; + const uint64_t ne22 = use_src2 ? src2->ne[2] : 0; + const uint64_t ne23 = use_src2 ? src2->ne[3] : 0; + const uint64_t ne2 = ne20 * ne21; + + const uint64_t ned0 = dst->ne[0]; + const uint64_t ned1 = dst->ne[1]; + const uint64_t ned2 = dst->ne[2]; + const uint64_t ned3 = dst->ne[3]; + const uint64_t ned = ned0 * ned1; + + init_pushconst_fastdiv(pc); + + vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, src2, dst, op); + + if (pipeline == nullptr) { + std::cerr << "ggml_vulkan: Error: Missing op: " << ggml_op_name(op) << " for " << ggml_type_name(src0->type); + if (src1 != nullptr) { + std::cerr << " and " << ggml_type_name(src1->type); + } + std::cerr << " to " << ggml_type_name(dst->type) << std::endl; + GGML_ABORT("fatal error"); + } + + if (dryrun) { + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + return; + } + + const bool op_supports_incontiguous = ggml_vk_op_supports_incontiguous(op); + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + ggml_backend_vk_buffer_context * src1_buf_ctx = use_src1 ? (ggml_backend_vk_buffer_context *)src1->buffer->context : nullptr; + ggml_backend_vk_buffer_context * src2_buf_ctx = use_src2 ? (ggml_backend_vk_buffer_context *)src2->buffer->context : nullptr; + + vk_buffer d_X = nullptr; + size_t x_buf_offset = 0; + vk_buffer d_Y = nullptr; + size_t y_buf_offset = 0; + vk_buffer d_Z = nullptr; + size_t z_buf_offset = 0; + + bool src0_uma = false; + bool src1_uma = false; + bool src2_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, src0->data, d_X, x_buf_offset); + src0_uma = d_X != nullptr; + if (use_src1) { + ggml_vk_host_get(ctx->device, src1->data, d_Y, y_buf_offset); + src1_uma = d_Y != nullptr; + } + if (use_src2) { + ggml_vk_host_get(ctx->device, src2->data, d_Z, z_buf_offset); + src2_uma = d_Z != nullptr; + } + } + + uint64_t x_sz = ggml_type_size(src0->type)/ggml_blck_size(src0->type) * ne0; + uint64_t y_sz = use_src1 ? ggml_type_size(src1->type) * ne1 : 0; + uint64_t z_sz = use_src2 ? ggml_type_size(src2->type) * ne2 : 0; + uint64_t d_sz = ggml_type_size(dst->type) * ned; + + vk_buffer d_D = dst_buf_ctx->dev_buffer; + + // Workaround for tiny tensor inputs on ROPE + if (op == GGML_OP_ROPE && use_src1 && y_sz > d_D->size) { + y_sz = VK_WHOLE_SIZE; + } + + GGML_ASSERT(d_D != nullptr); + uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + if(!src0_uma) { + d_X = src0_buf_ctx->dev_buffer; + x_buf_offset = vk_tensor_offset(src0) + src0->view_offs; + GGML_ASSERT(d_X != nullptr); + } + if (use_src1 && !src1_uma) { + d_Y = src1_buf_ctx->dev_buffer; + y_buf_offset = vk_tensor_offset(src1) + src1->view_offs; + GGML_ASSERT(d_Y != nullptr); + } + if (use_src2 && !src2_uma) { + d_Z = src2_buf_ctx->dev_buffer; + z_buf_offset = vk_tensor_offset(src2) + src2->view_offs; + GGML_ASSERT(d_Z != nullptr); + } + // Compute misalignment offset for descriptors and store it in in push constants, then align the descriptor offsets. + init_pushconst_tensor_offsets(ctx, pc, src0, src1, src2, dst); + x_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); + y_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); + z_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); + d_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); + + if (op_supports_incontiguous) { + x_sz = ggml_nbytes(src0); + y_sz = use_src1 ? ggml_nbytes(src1) : 0; + z_sz = use_src2 ? ggml_nbytes(src2) : 0; + d_sz = ggml_nbytes(dst); + + if (x_buf_offset + x_sz >= d_X->size) { + x_sz = VK_WHOLE_SIZE; + } + if (use_src1 && y_buf_offset + y_sz >= d_Y->size) { + y_sz = VK_WHOLE_SIZE; + } + if (use_src2 && z_buf_offset + z_sz >= d_Z->size) { + z_sz = VK_WHOLE_SIZE; + } + if (d_buf_offset + d_sz >= d_D->size) { + d_sz = VK_WHOLE_SIZE; + } + } + + std::array elements; + + // Single call if dimension 2 is contiguous + GGML_ASSERT(op_supports_incontiguous || (ggml_is_contiguous(src0) && (src1 == nullptr || ggml_is_contiguous(src1)))); + + switch (op) { + case GGML_OP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_SOFT_MAX: + case GGML_OP_SUM_ROWS: + { + const uint32_t nr = ggml_nrows(src0); + if (nr > 262144) { + elements = { 512, 512, CEIL_DIV(nr, 262144) }; + } else if (nr > 512) { + elements = { 512, CEIL_DIV(nr, 512), 1 }; + } else { + elements = { nr, 1, 1 }; + } + } break; + case GGML_OP_GROUP_NORM: + { + const uint32_t num_groups = dst->op_params[0]; + elements = { num_groups * (uint32_t)src0->ne[3], 1, 1 }; + } break; + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_ROPE: + elements = { (uint32_t)ggml_nrows(src0), (uint32_t)ne00, 1 }; + break; + case GGML_OP_GET_ROWS: + elements = { (uint32_t)ne00, (uint32_t)ne10, (uint32_t)(ne11 * ne12) }; + break; + case GGML_OP_ARGSORT: + elements = { (uint32_t)ne00, (uint32_t)ggml_nrows(src0), 1 }; + break; + case GGML_OP_IM2COL: + { + const bool is_2D = dst->op_params[6] == 1; + + const uint32_t IC = src1->ne[is_2D ? 2 : 1]; + + const uint32_t KH = is_2D ? src0->ne[1] : 1; + const uint32_t KW = src0->ne[0]; + + const uint32_t OH = is_2D ? dst->ne[2] : 1; + const uint32_t OW = dst->ne[1]; + + const uint32_t batch = src1->ne[is_2D ? 3 : 2]; + + elements = { OW * KW * KH, OH, batch * IC }; + } break; + case GGML_OP_TIMESTEP_EMBEDDING: + { + const uint32_t dim = dst->op_params[0]; + uint32_t half_ceil = (dim + 1) / 2; + elements = { half_ceil, (uint32_t)src0->ne[0], 1 }; + } break; + case GGML_OP_POOL_2D: + { + const uint32_t N = dst->ne[3]; + const uint32_t OC = dst->ne[2]; + const uint32_t OH = dst->ne[1]; + const uint32_t OW = dst->ne[0]; + elements = { N * OC * OH * OW, 1, 1}; + } break; + case GGML_OP_ADD: + case GGML_OP_DIV: + case GGML_OP_MUL: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_REPEAT: + case GGML_OP_CPY: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_UNARY: + { + const uint32_t ne = ggml_nelements(dst); + if (ne > 262144) { + elements = { 512, 512, CEIL_DIV(ne, 262144) }; + } else if (ne > 512) { + elements = { 512, CEIL_DIV(ne, 512), 1 }; + } else { + elements = { ne, 1, 1 }; + } + } break; + default: + elements = { (uint32_t)ggml_nelements(src0), 1, 1 }; + break; + } + + if (!op_supports_incontiguous) { + if (x_sz != VK_WHOLE_SIZE) { + x_sz *= ne02 * ne03; + } + if (use_src1 && y_sz != VK_WHOLE_SIZE) { + y_sz *= ne12 * ne13; + } + if (use_src2 && z_sz != VK_WHOLE_SIZE) { + z_sz *= ne22 * ne23; + } + if (d_sz != VK_WHOLE_SIZE) { + d_sz *= ned2 * ned3; + } + } + + if (op == GGML_OP_SOFT_MAX) { + // Empty src1 is possible in soft_max, but the shader needs a buffer + vk_subbuffer subbuf_y; + if (use_src1) { + subbuf_y = { d_Y, y_buf_offset, y_sz }; + } else { + subbuf_y = { d_X, 0, x_sz }; + } + + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, subbuf_y, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } else if (op == GGML_OP_ROPE) { + // Empty src2 is possible in rope, but the shader needs a buffer + vk_subbuffer subbuf_z; + if (use_src2) { + subbuf_z = { d_Z, z_buf_offset, z_sz }; + } else { + subbuf_z = { d_X, 0, x_sz }; + } + + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, subbuf_z, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } else if (op == GGML_OP_IM2COL) { + // im2col uses only src1 and dst buffers + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } else if (use_src2) { + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_Z, z_buf_offset, z_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } else if (use_src1) { + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } else { + ggml_vk_sync_buffers(subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); + } +} + +static void ggml_vk_get_rows(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_GET_ROWS, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, 0, + }, dryrun); +} + +static void ggml_vk_acc(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + int nb1 = dst->op_params[0] / 4; // 4 bytes of float32 + int nb2 = dst->op_params[1] / 4; // 4 bytes of float32 + // int nb3 = dst->op_params[2] / 4; // 4 bytes of float32 - unused + int offset = dst->op_params[3] / 4; // offset in bytes + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_ACC, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)nb1, (uint32_t)nb2, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t)nb1, (uint32_t)nb2, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, offset, + }, dryrun); +} + +static void ggml_vk_add(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_ADD, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, 0, + }, dryrun); +} + +static void ggml_vk_mul(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_MUL, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, 0, + }, dryrun); +} + +static void ggml_vk_div(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_DIV, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, 0, + }, dryrun); +} + +static void ggml_vk_op_f32_rwkv6(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, const vk_op_rwkv_wkv6_push_constants&& pc, bool dryrun = false) { + const ggml_tensor * k = dst->src[0]; + const ggml_tensor * v = dst->src[1]; + const ggml_tensor * r = dst->src[2]; + const ggml_tensor * tf = dst->src[3]; + const ggml_tensor * td = dst->src[4]; + const ggml_tensor * state = dst->src[5]; + + GGML_ASSERT(!ggml_is_quantized(k->type)); + GGML_ASSERT(!ggml_is_quantized(v->type)); + GGML_ASSERT(!ggml_is_quantized(r->type)); + GGML_ASSERT(!ggml_is_quantized(tf->type)); + GGML_ASSERT(!ggml_is_quantized(td->type)); + GGML_ASSERT(!ggml_is_quantized(state->type)); + GGML_ASSERT(dst->buffer != nullptr); + + vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, k, v, r, dst, GGML_OP_RWKV_WKV6); + GGML_ASSERT(pipeline != nullptr); + + if (dryrun) { + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + return; + } + + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + ggml_backend_vk_buffer_context * k_buf_ctx = (ggml_backend_vk_buffer_context *)k->buffer->context; + ggml_backend_vk_buffer_context * v_buf_ctx = (ggml_backend_vk_buffer_context *)v->buffer->context; + ggml_backend_vk_buffer_context * r_buf_ctx = (ggml_backend_vk_buffer_context *)r->buffer->context; + ggml_backend_vk_buffer_context * tf_buf_ctx = (ggml_backend_vk_buffer_context *)tf->buffer->context; + ggml_backend_vk_buffer_context * td_buf_ctx = (ggml_backend_vk_buffer_context *)td->buffer->context; + ggml_backend_vk_buffer_context * state_buf_ctx = (ggml_backend_vk_buffer_context *)state->buffer->context; + + ggml_vk_sync_buffers(subctx); + + vk_buffer d_D = nullptr, d_K = nullptr, d_V = nullptr, d_R = nullptr, d_TF = nullptr, d_TD = nullptr, d_State = nullptr; + size_t k_offset = 0, v_offset = 0, r_offset = 0, tf_offset = 0, td_offset = 0, state_offset = 0, dst_offset = 0; + bool K_uma = false, V_uma = false, R_uma = false, TF_uma = false, TD_uma = false, STATE_uma = false, DST_uma = false; + + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, k->data, d_K, k_offset); + ggml_vk_host_get(ctx->device, v->data, d_V, v_offset); + ggml_vk_host_get(ctx->device, r->data, d_R, r_offset); + ggml_vk_host_get(ctx->device, tf->data, d_TF, tf_offset); + ggml_vk_host_get(ctx->device, td->data, d_TD, td_offset); + ggml_vk_host_get(ctx->device, state->data, d_State, state_offset); + ggml_vk_host_get(ctx->device, dst->data, d_D, dst_offset); + + K_uma = d_K != nullptr; + V_uma = d_V != nullptr; + R_uma = d_R != nullptr; + TF_uma = d_TF != nullptr; + TD_uma = d_TD != nullptr; + STATE_uma = d_State != nullptr; + DST_uma = d_D != nullptr; + } + + if (!K_uma) { + d_K = k_buf_ctx->dev_buffer; + k_offset = vk_tensor_offset(k) + k->view_offs; + } + if (!V_uma) { + d_V = v_buf_ctx->dev_buffer; + v_offset = vk_tensor_offset(v) + v->view_offs; + } + if (!R_uma) { + d_R = r_buf_ctx->dev_buffer; + r_offset = vk_tensor_offset(r) + r->view_offs; + } + if (!TF_uma) { + d_TF = tf_buf_ctx->dev_buffer; + tf_offset = vk_tensor_offset(tf) + tf->view_offs; + } + if (!TD_uma) { + d_TD = td_buf_ctx->dev_buffer; + td_offset = vk_tensor_offset(td) + td->view_offs; + } + if (!STATE_uma) { + d_State = state_buf_ctx->dev_buffer; + state_offset = vk_tensor_offset(state) + state->view_offs; + } + if (!DST_uma) { + d_D = dst_buf_ctx->dev_buffer; + dst_offset = vk_tensor_offset(dst) + dst->view_offs; + } + + const uint64_t k_size = ggml_nbytes(k); + const uint64_t v_size = ggml_nbytes(v); + const uint64_t r_size = ggml_nbytes(r); + const uint64_t tf_size = ggml_nbytes(tf); + const uint64_t td_size = ggml_nbytes(td); + const uint64_t state_size = ggml_nbytes(state); + const uint64_t dst_size = ggml_nbytes(dst); + + std::array elements = { + (uint32_t)(pc.B * pc.H), + 1, + 1 + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { + vk_subbuffer{ d_K, k_offset, k_size }, + vk_subbuffer{ d_V, v_offset, v_size }, + vk_subbuffer{ d_R, r_offset, r_size }, + vk_subbuffer{ d_TF, tf_offset, tf_size }, + vk_subbuffer{ d_TD, td_offset, td_size }, + vk_subbuffer{ d_State, state_offset, state_size }, + vk_subbuffer{ d_D, dst_offset, dst_size } + }, sizeof(vk_op_rwkv_wkv6_push_constants), &pc, elements); +} + +static void ggml_vk_rwkv_wkv6(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, bool dryrun = false) { + const size_t seq_length = dst->src[0]->ne[3]; + const size_t n_embed = dst->ne[0]; + const size_t n_heads = dst->src[0]->ne[2]; + const size_t n_seqs = dst->src[5]->ne[1]; + + ggml_vk_op_f32_rwkv6( + ctx, subctx, dst, + { + (uint32_t)n_seqs, + (uint32_t)seq_length, + (uint32_t)n_embed, + (uint32_t)n_heads, + }, + dryrun + ); +} + +static void ggml_vk_concat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + int * op_params = (int *)dst->op_params; + + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_CONCAT, { + (uint32_t)ggml_nelements(dst), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, op_params[0], + }, dryrun); +} + +static void ggml_vk_upscale(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + + const float sf0 = (float)dst->ne[0] / src0->ne[0]; + const float sf1 = (float)dst->ne[1] / src0->ne[1]; + const float sf2 = (float)dst->ne[2] / src0->ne[2]; + const float sf3 = (float)dst->ne[3] / src0->ne[3]; + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_UPSCALE, { + (uint32_t)ggml_nelements(dst), 0, 0, + (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)dst->ne[0], (uint32_t)dst->ne[1], (uint32_t)dst->ne[2],(uint32_t)dst->ne[3], + sf0, sf1, sf2, sf3, + }, dryrun); +} + +static void ggml_vk_scale(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + float * op_params = (float *)dst->op_params; + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SCALE, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + op_params[0], 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_sqr(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SQR, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_sin(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SIN, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_cos(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_COS, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_clamp(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + float * op_params = (float *)dst->op_params; + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_CLAMP, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + op_params[0], op_params[1], + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_pad(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_PAD, { + (uint32_t)ggml_nelements(dst), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_repeat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_REPEAT, { + (uint32_t)ggml_nelements(dst), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_cpy(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_CPY, { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }, dryrun); +} + +static void ggml_vk_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + float * op_params = (float *)dst->op_params; + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_NORM, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0], 0.0f }, dryrun); +} + +static void ggml_vk_group_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const int * int_op_params = (const int *)dst->op_params; + const float * float_op_params = (const float *)dst->op_params; + + const uint32_t num_groups = int_op_params[0]; + const float eps = float_op_params[1]; + const uint32_t group_size = src0->ne[0] * src0->ne[1] * ((src0->ne[2] + num_groups - 1) / num_groups); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_GROUP_NORM, { group_size, 0, eps, 0.0f }, dryrun); +} + +static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + float * op_params = (float *)dst->op_params; + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_RMS_NORM, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0], 0.0f }, dryrun); +} + +static void ggml_vk_unary(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_UNARY, { (uint32_t)ggml_nelements(src0), 0, 0.0f, 0.0f }, dryrun); +} + +static void ggml_vk_diag_mask_inf(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + int32_t * op_params = (int32_t *)dst->op_params; + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_DIAG_MASK_INF, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0] }, dryrun); +} + +static void ggml_vk_soft_max(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + float * op_params = (float *)dst->op_params; + + float scale = op_params[0]; + float max_bias = op_params[1]; + + const uint32_t ncols = (uint32_t)src0->ne[0]; + const uint32_t nrows_x = (uint32_t)ggml_nrows(src0); + const uint32_t nrows_y = (uint32_t)src0->ne[1]; + + const uint32_t n_head_kv = nrows_x/nrows_y; + const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head_kv)); + + const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); + const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_SOFT_MAX, { + ncols, + src1 != nullptr ? nrows_y : (uint32_t)0, + scale, max_bias, + m0, m1, + n_head_log2, + nrows_x, + }, dryrun); +} + +static void ggml_vk_rope(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, bool dryrun = false) { + const int n_dims = ((int32_t *) dst->op_params)[1]; + // const int mode = ((int32_t *) dst->op_params)[2]; + // const int n_ctx = ((int32_t *) dst->op_params)[3]; + const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; + const float freq_base = ((float *) dst->op_params)[5]; + const float freq_scale = ((float *) dst->op_params)[6]; + const float ext_factor = ((float *) dst->op_params)[7]; + const float attn_factor = ((float *) dst->op_params)[8]; + const float beta_fast = ((float *) dst->op_params)[9]; + const float beta_slow = ((float *) dst->op_params)[10]; + + float corr_dims[2]; + ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); + + const float theta_scale = powf(freq_base, -2.0f/n_dims); + + ggml_vk_op_f32(ctx, subctx, src0, src1, src2, dst, GGML_OP_ROPE, { + (uint32_t)src0->ne[0], (uint32_t)n_dims, freq_scale, (uint32_t)src0->ne[1], + freq_base, ext_factor, attn_factor, {corr_dims[0], corr_dims[1]}, theta_scale, + src2 != nullptr, + }, dryrun); +} + +static void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + int32_t * op_params = (int32_t *)dst->op_params; + + uint32_t ncols = src0->ne[0]; + + uint32_t ncols_pad = 1; + while (ncols_pad < ncols) { + ncols_pad *= 2; + } + + GGML_ASSERT(ncols_pad <= 1024); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_ARGSORT, { + ncols, + ncols_pad, + op_params[0], + }, dryrun); +} + +static void ggml_vk_sum_rows(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SUM_ROWS, { (uint32_t)src0->ne[0], 0, 0.0f, 0.0f }, dryrun); +} + +static void ggml_vk_im2col(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { + const int32_t s0 = dst->op_params[0]; + const int32_t s1 = dst->op_params[1]; + const int32_t p0 = dst->op_params[2]; + const int32_t p1 = dst->op_params[3]; + const int32_t d0 = dst->op_params[4]; + const int32_t d1 = dst->op_params[5]; + + const bool is_2D = dst->op_params[6] == 1; + + const uint32_t IC = src1->ne[is_2D ? 2 : 1]; + const uint32_t IH = is_2D ? src1->ne[1] : 1; + const uint32_t IW = src1->ne[0]; + + const uint32_t KH = is_2D ? src0->ne[1] : 1; + const uint32_t KW = src0->ne[0]; + + const uint32_t OH = is_2D ? dst->ne[2] : 1; + const uint32_t OW = dst->ne[1]; + + const uint32_t offset_delta = src1->nb[is_2D ? 2 : 1] / 4; // nb is byte offset, src is type float32 + const uint32_t batch_offset = src1->nb[is_2D ? 3 : 2] / 4; // nb is byte offset, src is type float32 + + const uint32_t pelements = OW * KW * KH; + + ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_IM2COL, { + batch_offset, offset_delta, + IC, IW, IH, OW, OH, KW, KH, + pelements, + IC * KH * KW, + s0, s1, p0, p1, d0, d1, + }, dryrun); +} + +static void ggml_vk_timestep_embedding(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const uint32_t dim = dst->op_params[0]; + const uint32_t max_period = dst->op_params[1]; + const uint32_t nb1 = dst->nb[1] / ggml_type_size(dst->type); + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_TIMESTEP_EMBEDDING, { + nb1, dim, max_period, + }, dryrun); +} + +static void ggml_vk_pool_2d(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + uint32_t op = static_cast(dst->op_params[0]); + const int32_t k1 = dst->op_params[1]; + const int32_t k0 = dst->op_params[2]; + const int32_t s1 = dst->op_params[3]; + const int32_t s0 = dst->op_params[4]; + const int32_t p1 = dst->op_params[5]; + const int32_t p0 = dst->op_params[6]; + + const uint32_t IH = src0->ne[1]; + const uint32_t IW = src0->ne[0]; + + const uint32_t N = dst->ne[3]; + + const uint32_t OC = dst->ne[2]; + const uint32_t OH = dst->ne[1]; + const uint32_t OW = dst->ne[0]; + + const uint32_t parallel_elements = N * OC * OH * OW; + + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_POOL_2D, { + IW, IH, OW, OH, OC, + parallel_elements, + op, + k0, k1, s0, s1, p0, p1, + }, dryrun); +} + +static void ggml_vk_leaky_relu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { + const float * op_params = (const float *)dst->op_params; + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_LEAKY_RELU, { (uint32_t)ggml_nelements(src0), 0, op_params[0], 0.0f }, dryrun); +} + +#ifdef GGML_VULKAN_RUN_TESTS +static void ggml_vk_print_matrix_area(const void * data, ggml_type type, int ne0, int ne1, int i0, int i1, int i2) { + if (type != GGML_TYPE_F32 && type != GGML_TYPE_F16) { + return; + } + i0 = std::max(i0, 5); + i1 = std::max(i1, 5); + i2 = std::max(i2, 0); + fprintf(stderr, " "); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + fprintf(stderr, "%7d ", idx1); + } + fprintf(stderr, "\n"); + for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { + fprintf(stderr, "%7d: ", idx0); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + if (idx0 >= 0 && idx0 < ne0 && idx1 >= 0 && idx1 < ne1) { + float val; + if (type == GGML_TYPE_F32) { + val = *((const float *) data + i2*ne1*ne0 + idx1*ne0 + idx0); + } else if (type == GGML_TYPE_F16) { + val = ggml_fp16_to_fp32(*((const ggml_fp16_t *) data + i2*ne1*ne0 + idx1*ne0 + idx0)); + } else { + GGML_ABORT("fatal error"); + } + fprintf(stderr, "% 7.2f ", val); + } else { + fprintf(stderr, " "); + } + } + fprintf(stderr, "\n"); + } +} + +template +static void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, int split_k, int shader_size) { + VK_LOG_DEBUG("ggml_vk_test_matmul(" << m << ", " << n << ", " << k << ", " << batch << ", " << num_it << ", " << split_k << ", " << shader_size << ")"); + const size_t x_ne = m * k * batch; + const size_t y_ne = k * n * batch; + const size_t d_ne = m * n * batch; + + vk_pipeline p; + std::string shname; + if (shader_size == 0) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->a_s; + shname = "F32_ALIGNED_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->a_s; + shname = "F32_F16_ALIGNED_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_s; + shname = "F16_F32_ALIGNED_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->a_s; + shname = "F16_ALIGNED_S"; + } else { + GGML_ABORT("fatal error"); + } + } else if (shader_size == 1) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->a_m; + shname = "F32_ALIGNED_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->a_m; + shname = "F32_F16_ALIGNED_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_m; + shname = "F16_F32_ALIGNED_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->a_m; + shname = "F16_ALIGNED_M"; + } else { + GGML_ABORT("fatal error"); + } + } else if (shader_size == 2) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->a_l; + shname = "F32_ALIGNED_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->a_l; + shname = "F32_F16_ALIGNED_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_l; + shname = "F16_F32_ALIGNED_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->a_l; + shname = "F16_ALIGNED_L"; + } else { + GGML_ABORT("fatal error"); + } + } else { + GGML_ASSERT(0); + } + + const size_t kpad = ggml_vk_align_size(k, p->align); + + if (k != kpad) { + if (shader_size == 0) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->s; + shname = "F32_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->s; + shname = "F32_F16_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->s; + shname = "F16_F32_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->s; + shname = "F16_S"; + } + } else if (shader_size == 1) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->m; + shname = "F32_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->m; + shname = "F32_F16_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->m; + shname = "F16_F32_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->m; + shname = "F16_M"; + } + } else if (shader_size == 2) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->l; + shname = "F32_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->l; + shname = "F32_F16_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->l; + shname = "F16_F32_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->l; + shname = "F16_L"; + } + } + } + + ggml_pipeline_request_descriptor_sets(ctx->device, p, num_it); + if (split_k > 1) { + ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, num_it); + + if (ctx->prealloc_split_k == nullptr || ctx->prealloc_split_k->size < sizeof(float) * d_ne * split_k) { + // Resize buffer + if (ctx->prealloc_split_k != nullptr) { + ggml_vk_destroy_buffer(ctx->prealloc_split_k); + } + ctx->prealloc_split_k = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne * split_k, vk::MemoryPropertyFlagBits::eDeviceLocal); + } + } + + ggml_pipeline_allocate_descriptor_sets(ctx->device); + + vk_buffer d_X = ggml_vk_create_buffer_check(ctx->device, sizeof(X_TYPE) * x_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); + vk_buffer d_Y = ggml_vk_create_buffer_check(ctx->device, sizeof(Y_TYPE) * y_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); + vk_buffer d_D = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); + + X_TYPE* x = (X_TYPE *) malloc(sizeof(X_TYPE) * x_ne); + Y_TYPE* y = (Y_TYPE *) malloc(sizeof(Y_TYPE) * y_ne); + float* d = (float *) malloc(sizeof(float) * d_ne); + + for (size_t i = 0; i < x_ne; i++) { + if (std::is_same()) { + x[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; + // x[i] = 1.0f; + // x[i] = i + 1; + // x[i] = (i % k == i / k) ? 1.0f : 0.0f; + } else if (std::is_same()) { + x[i] = ggml_fp32_to_fp16((rand() / (float)RAND_MAX) * 2.0f - 1.0f); + // x[i] = ggml_fp32_to_fp16(1.0f); + // x[i] = ggml_fp32_to_fp16(i + 1); + // x[i] = ggml_fp32_to_fp16((i % k == i / k) ? 1.0f : 0.0f); + } else { + GGML_ABORT("fatal error"); + } + } + for (size_t i = 0; i < y_ne; i++) { + if (std::is_same()) { + y[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; + // y[i] = (i % k == i / k) ? 1.0f : 0.0f; + // y[i] = i + 1; + } else if (std::is_same()) { + y[i] = ggml_fp32_to_fp16((rand() / (float)RAND_MAX) * 2.0f - 1.0f); + // y[i] = ggml_fp32_to_fp16((i % k == i / k) ? 1.0f : 0.0f); + // y[i] = ggml_fp32_to_fp16(i + 1); + } else { + GGML_ABORT("fatal error"); + } + } + + ggml_vk_buffer_write(d_X, 0, x, sizeof(X_TYPE) * k * m * batch); + ggml_vk_buffer_write(d_Y, 0, y, sizeof(Y_TYPE) * k * n * batch); + + vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); + ggml_vk_ctx_begin(ctx->device, subctx); + for (size_t i = 0; i < num_it; i++) { + ggml_vk_matmul( + ctx, subctx, p, ggml_vk_subbuffer(d_X), ggml_vk_subbuffer(d_Y), ggml_vk_subbuffer(d_D), ggml_vk_subbuffer(ctx->prealloc_split_k), + m, n, k, + k, k, m, k*m, k*n, m*n, + split_k, batch, batch, batch, 1, 1 + ); + } + ggml_vk_ctx_end(subctx); + + auto begin = std::chrono::high_resolution_clock::now(); + ggml_vk_submit(subctx, ctx->fence); + VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_matmul waitForFences"); + ctx->device->device.resetFences({ ctx->fence }); + + auto end = std::chrono::high_resolution_clock::now(); + double time = std::chrono::duration_cast(end-begin).count() / 1000.0; + + // copy dst to host + ggml_vk_buffer_read(d_D, 0, d, sizeof(float) * d_ne); + + float * d_chk = (float *) malloc(sizeof(float) * d_ne); + + ggml_init_params iparams = { + /*.mem_size =*/ 1024*1024*1024, + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + + ggml_context * ggml_ctx = ggml_init(iparams); + + ggml_type src0_type; + ggml_type src1_type; + + if (std::is_same()) { + src0_type = GGML_TYPE_F32; + } else if (std::is_same()) { + src0_type = GGML_TYPE_F16; + } else { + GGML_ABORT("fatal error"); + } + if (std::is_same()) { + src1_type = GGML_TYPE_F32; + } else if (std::is_same()) { + src1_type = GGML_TYPE_F16; + } else { + GGML_ABORT("fatal error"); + } + + ggml_tensor * src0_ggml = ggml_new_tensor_3d(ggml_ctx, src0_type, k, m, batch); + ggml_tensor * src1_ggml = ggml_new_tensor_3d(ggml_ctx, src1_type, k, n, batch); + ggml_tensor * tensor_ggml = ggml_mul_mat(ggml_ctx, src0_ggml, src1_ggml); + + src0_ggml->data = x; + src1_ggml->data = y; + tensor_ggml->data = d_chk; + + ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); + ggml_build_forward_expand(cgraph, tensor_ggml); + + ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 1); + + ggml_free(ggml_ctx); + + double avg_err = 0.0; + int first_err_n = -1; + int first_err_m = -1; + int first_err_b = -1; + + for (size_t i = 0; i < m*n*batch; i++) { + double err = std::fabs(d[i] - d_chk[i]); + avg_err += err; + + if ((err > 0.05f || std::isnan(err)) && first_err_n == -1) { + first_err_b = i / (m * n); + first_err_n = (i % (m * n)) / m; + first_err_m = (i % (m * n)) % m; + } + } + + avg_err /= m * n; + + double tflops = 2.0*m*n*k*batch*num_it / (time / 1000.0) / (1000.0*1000.0*1000.0*1000.0); + + std::cerr << "TEST " << shname << " m=" << m << " n=" << n << " k=" << k << " batch=" << batch << " split_k=" << split_k << " matmul " << time / num_it << "ms " << tflops << " TFLOPS avg_err=" << avg_err << std::endl; + + if (avg_err > 0.1 || std::isnan(avg_err)) { + std::cerr << "m = " << first_err_m << " n = " << first_err_n << " b = " << first_err_b << std::endl; + std::cerr << "Actual result: " << std::endl << std::endl; + ggml_vk_print_matrix_area(d, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + std::cerr << "Expected result: " << std::endl << std::endl; + ggml_vk_print_matrix_area(d_chk, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + if (split_k > 1) { + float * split_k_buf = (float *) malloc(sizeof(float) * d_ne * split_k); + ggml_vk_buffer_read(ctx->prealloc_split_k, 0, split_k_buf, sizeof(float) * d_ne * split_k); + + std::cerr << "d_buf0: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf1: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf2: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + 2 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf3: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + 3 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + free(split_k_buf); + } + } + + free(d_chk); + + ggml_vk_queue_cleanup(ctx->device, ctx->device->transfer_queue); + ggml_vk_queue_cleanup(ctx->device, ctx->device->compute_queue); + + ggml_vk_destroy_buffer(d_X); + ggml_vk_destroy_buffer(d_Y); + ggml_vk_destroy_buffer(d_D); + + ggml_pipeline_cleanup(p); + ggml_pipeline_cleanup(ctx->device->pipeline_matmul_split_k_reduce); + + free(x); + free(y); + free(d); +} + +static void ggml_vk_print_tensor_area(const ggml_tensor * tensor, int i0, int i1, int i2, int i3) { + if (tensor->type != GGML_TYPE_F32 && tensor->type != GGML_TYPE_F16) { + return; + } + i0 = std::max(i0, 5); + i1 = std::max(i1, 5); + i2 = std::max(i2, 0); + i3 = std::max(i3, 0); + fprintf(stderr, " "); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + fprintf(stderr, "%7d ", idx1); + } + fprintf(stderr, "\n"); + for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { + fprintf(stderr, "%7d: ", idx0); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + if (idx0 >= 0 && idx0 < tensor->ne[0] && idx1 >= 0 && idx1 < tensor->ne[1] && i2 >= 0 && i2 < tensor->ne[2] && i3 >= 0 && i3 < tensor->ne[3]) { + float val; + if (tensor->type == GGML_TYPE_F32) { + val = *(float *) ((char *) tensor->data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); + } else if (tensor->type == GGML_TYPE_F16) { + val = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) tensor->data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0])); + } else { + GGML_ABORT("fatal error"); + } + fprintf(stderr, "% 7.2f ", val); + } else { + fprintf(stderr, " "); + } + } + fprintf(stderr, "\n"); + } +} + +static void ggml_vk_quantize_data(const float * from, void * to, size_t ne, ggml_type quant) { + ggml_quantize_chunk(quant, from, to, 0, 1, ne, nullptr); +} + +static void ggml_vk_dequantize_data(const void * from, float * to, size_t ne, ggml_type quant) { + if (quant == GGML_TYPE_F32) { + memcpy(to, from, sizeof(float) * ne); + return; + } + + const auto * tt = ggml_get_type_traits(quant); + + ggml_to_float_t dequant_fn = tt->to_float; + + dequant_fn(from, to, ne); +} + +static void ggml_vk_test_dequant(ggml_backend_vk_context * ctx, size_t ne, ggml_type quant) { + VK_LOG_DEBUG("ggml_vk_test_dequant(" << ne << ")"); + const size_t x_sz = sizeof(float) * ne; + const size_t x_sz_f16 = sizeof(ggml_fp16_t) * ne; + const size_t qx_sz = ne * ggml_type_size(quant)/ggml_blck_size(quant); + float * x = (float *) malloc(x_sz); + void * qx = malloc(qx_sz); + vk_buffer qx_buf = ggml_vk_create_buffer_check(ctx->device, qx_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); + vk_buffer x_buf = ggml_vk_create_buffer_check(ctx->device, x_sz_f16, vk::MemoryPropertyFlagBits::eDeviceLocal); + float * x_ref = (float *) malloc(x_sz); + ggml_fp16_t * x_chk = (ggml_fp16_t *) malloc(x_sz_f16); + + for (size_t i = 0; i < ne; i++) { + x[i] = rand() / (float)RAND_MAX; + } + + vk_pipeline p = ggml_vk_get_to_fp16(ctx, quant); + + ggml_vk_quantize_data(x, qx, ne, quant); + ggml_vk_dequantize_data(qx, x_ref, ne, quant); + + ggml_pipeline_request_descriptor_sets(ctx->device, p, 1); + + ggml_pipeline_allocate_descriptor_sets(ctx->device); + + ggml_vk_buffer_write(qx_buf, 0, qx, qx_sz); + + vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); + ggml_vk_ctx_begin(ctx->device, subctx); + const std::vector pc = { 1, (uint32_t)ne, (uint32_t)ne, (uint32_t)ne, (uint32_t)ne }; + ggml_vk_dispatch_pipeline(ctx, subctx, p, { vk_subbuffer{ qx_buf, 0, qx_sz }, vk_subbuffer{ x_buf, 0, x_sz_f16 } }, pc.size() * sizeof(int), pc.data(), { (uint32_t)ne, 1, 1}); + ggml_vk_ctx_end(subctx); + + auto begin = std::chrono::high_resolution_clock::now(); + + ggml_vk_submit(subctx, ctx->fence); + VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences"); + ctx->device->device.resetFences({ ctx->fence }); + + auto end = std::chrono::high_resolution_clock::now(); + + double ms_dequant = std::chrono::duration_cast(end-begin).count() / 1000.0; + ggml_vk_buffer_read(x_buf, 0, x_chk, x_sz_f16); + + int first_err = -1; + + double avg_err = 0.0; + for (size_t i = 0; i < ne; i++) { + double error = std::fabs(x_ref[i] - ggml_fp16_to_fp32(x_chk[i])); + avg_err += error; + + if (first_err < 0 && error > 0.05) { + first_err = i; + } + } + + avg_err /= ne; + + std::cerr << "TEST DEQUANT " << ggml_type_name(quant) << " time=" << ms_dequant << "ms avg_err=" << avg_err << std::endl; + + if (avg_err > 0.1) { + std::cerr << "first_error = " << first_err << std::endl; + std::cerr << "Actual result: " << std::endl << std::endl; + for (int i = std::max(0, first_err - 5); i < std::min((int)ne, first_err + 5); i++) { + std::cerr << ggml_fp16_to_fp32(x_chk[i]) << ", "; + } + std::cerr << std::endl << "Expected result: " << std::endl << std::endl; + for (int i = std::max(0, first_err - 5); i < std::min((int)ne, first_err + 5); i++) { + std::cerr << x_ref[i] << ", "; + } + std::cerr << std::endl; + } + + ggml_vk_destroy_buffer(x_buf); + ggml_vk_destroy_buffer(qx_buf); + + free(x); + free(qx); + free(x_ref); + free(x_chk); +} + +static void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, size_t split_k, size_t shader_size, ggml_type quant) { + VK_LOG_DEBUG("ggml_vk_test_dequant_matmul(" << m << ", " << n << ", " << k << ", " << batch << ", " << num_it << ", " << split_k << ", " << ggml_type_name(quant) << ")"); + const size_t x_ne = m * k * batch; + const size_t y_ne = k * n * batch; + const size_t d_ne = m * n * batch; + + vk_pipeline p; + std::string shname; + if (shader_size == 0) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_s : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_s; + shname = std::string(ggml_type_name(quant)) + "_ALIGNED_S"; + } else if (shader_size == 1) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_m : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_m; + shname = std::string(ggml_type_name(quant)) + "_ALIGNED_M"; + } else if (shader_size == 2) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_l : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_l; + shname = std::string(ggml_type_name(quant)) + "_ALIGNED_L"; + } else { + GGML_ASSERT(0); + } + + const size_t kpad = ggml_vk_align_size(k, p->align); + + if (k != kpad) { + if (shader_size == 0) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->s : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->s; + shname = std::string(ggml_type_name(quant)) + "_S"; + } else if (shader_size == 1) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->m : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->m; + shname = std::string(ggml_type_name(quant)) + "_M"; + } else if (shader_size == 2) { + p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->l : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->l; + shname = std::string(ggml_type_name(quant)) + "_L"; + } else { + GGML_ASSERT(0); + } + } + + const size_t x_sz = sizeof(float) * x_ne; + const size_t y_sz = sizeof(float) * y_ne; + const size_t qx_sz = x_ne * ggml_type_size(quant)/ggml_blck_size(quant); + const size_t d_sz = sizeof(float) * d_ne; + float * x = (float *) malloc(x_sz); + float * y = (float *) malloc(y_sz); + void * qx = malloc(qx_sz); + vk_buffer qx_buf = ggml_vk_create_buffer_check(ctx->device, qx_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); + vk_buffer y_buf = ggml_vk_create_buffer_check(ctx->device, y_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); + vk_buffer d_buf = ggml_vk_create_buffer_check(ctx->device, d_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); + float * d = (float *) malloc(d_sz); + float * d_chk = (float *) malloc(d_sz); + + for (size_t i = 0; i < x_ne; i++) { + x[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; + } + + ggml_vk_quantize_data(x, qx, x_ne, quant); + + for (size_t i = 0; i < y_ne; i++) { + // y[i] = rand() / (float)RAND_MAX; + y[i] = (i % k == i / k) ? 1.0f : 0.0f; + } + + ggml_pipeline_request_descriptor_sets(ctx->device, p, num_it); + if (split_k > 1) { + ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, num_it); + + if (ctx->prealloc_split_k == nullptr || ctx->prealloc_split_k->size < sizeof(float) * d_ne * split_k) { + // Resize buffer + if (ctx->prealloc_split_k != nullptr) { + ggml_vk_destroy_buffer(ctx->prealloc_split_k); + } + ctx->prealloc_split_k = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne * split_k, vk::MemoryPropertyFlagBits::eDeviceLocal); + } + } + + ggml_pipeline_allocate_descriptor_sets(ctx->device); + + ggml_vk_buffer_write(qx_buf, 0, qx, qx_sz); + ggml_vk_buffer_write(y_buf, 0, y, y_sz); + + vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); + ggml_vk_ctx_begin(ctx->device, subctx); + for (size_t i = 0; i < num_it; i++) { + ggml_vk_matmul( + ctx, subctx, p, ggml_vk_subbuffer(qx_buf), ggml_vk_subbuffer(y_buf), ggml_vk_subbuffer(d_buf), ggml_vk_subbuffer(ctx->prealloc_split_k), + m, n, k, + k, k, m, k*m, k*n, m*n, + split_k, batch, batch, batch, 1, 1 + ); + } + ggml_vk_ctx_end(subctx); + + auto begin = std::chrono::high_resolution_clock::now(); + + ggml_vk_submit(subctx, ctx->fence); + VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences"); + ctx->device->device.resetFences({ ctx->fence }); + + auto end = std::chrono::high_resolution_clock::now(); + + double time_ms = std::chrono::duration_cast(end-begin).count() / 1000.0; + ggml_vk_buffer_read(d_buf, 0, d, d_sz); + + ggml_init_params iparams = { + /*.mem_size =*/ 1024*1024*1024, + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + + ggml_context * ggml_ctx = ggml_init(iparams); + + ggml_tensor * src0_ggml = ggml_new_tensor_3d(ggml_ctx, quant, k, m, batch); + ggml_tensor * src1_ggml = ggml_new_tensor_3d(ggml_ctx, GGML_TYPE_F32, k, n, batch); + ggml_tensor * tensor_ggml = ggml_mul_mat(ggml_ctx, src0_ggml, src1_ggml); + + src0_ggml->data = qx; + src1_ggml->data = y; + tensor_ggml->data = d_chk; + + ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); + ggml_build_forward_expand(cgraph, tensor_ggml); + + ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 1); + + ggml_free(ggml_ctx); + + double avg_err = 0.0; + int first_err_n = -1; + int first_err_m = -1; + int first_err_b = -1; + + for (size_t i = 0; i < m*n*batch; i++) { + double err = std::fabs(d[i] - d_chk[i]); + avg_err += err; + + if ((err > 0.05f || std::isnan(err)) && first_err_n == -1) { + first_err_b = i / (m * n); + first_err_n = (i % (m * n)) / m; + first_err_m = (i % (m * n)) % m; + } + } + + avg_err /= m * n; + + double tflops = 2.0*m*n*k*batch*num_it / (time_ms / 1000.0) / (1000.0*1000.0*1000.0*1000.0); + + std::cerr << "TEST MMQ " << shname << " m=" << m << " n=" << n << " k=" << k << " batch=" << batch << " split_k=" << split_k << " matmul " << time_ms / num_it << "ms " << tflops << " TFLOPS avg_err=" << avg_err << std::endl; + + if (avg_err > 0.01 || std::isnan(avg_err)) { + std::cerr << "m = " << first_err_m << " n = " << first_err_n << " b = " << first_err_b << std::endl; + std::cerr << "Actual result: " << std::endl << std::endl; + ggml_vk_print_matrix_area(d, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + std::cerr << std::endl; + std::cerr << "Expected result: " << std::endl << std::endl; + ggml_vk_print_matrix_area(d_chk, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + if (split_k > 1) { + float * split_k_buf = (float *) malloc(sizeof(float) * d_ne * split_k); + ggml_vk_buffer_read(ctx->prealloc_split_k, 0, split_k_buf, sizeof(float) * d_ne * split_k); + + std::cerr << "d_buf0: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf1: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf2: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + 2 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + std::cerr << "d_buf3: " << std::endl << std::endl; + ggml_vk_print_matrix_area(split_k_buf + 3 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); + + free(split_k_buf); + } + } + + ggml_vk_destroy_buffer(qx_buf); + ggml_vk_destroy_buffer(y_buf); + ggml_vk_destroy_buffer(d_buf); + + free(x); + free(qx); + free(y); + free(d); + free(d_chk); +} +#endif + +static void ggml_vk_preallocate_buffers(ggml_backend_vk_context * ctx) { +#if defined(GGML_VULKAN_RUN_TESTS) + const std::vector vals { + 512, 512, 128, + 128, 512, 512, + 4096, 512, 4096, + 11008, 512, 4096, + 4096, 512, 11008, + 32000, 512, 4096, + 8, 8, 8, + 100, 46, 576, + 623, 111, 128, + 100, 46, 558, + 512, 1, 256, + 128, 110, 622, + 511, 511, 127, + 511, 511, 7, + 511, 511, 17, + 49, 49, 128, + 128, 49, 49, + 4096, 49, 4096, + }; + const size_t num_it = 100; + + for (size_t i = 0; i < vals.size(); i += 3) { + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2); + std::cerr << '\n'; + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2); + std::cerr << '\n'; + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1); + ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2); + std::cerr << '\n' << std::endl; + + if (vals[i + 2] % 32 == 0) { + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2, GGML_TYPE_Q4_0); + std::cerr << '\n'; + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2, GGML_TYPE_Q4_0); + std::cerr << '\n'; + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1, GGML_TYPE_Q4_0); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2, GGML_TYPE_Q4_0); + std::cerr << '\n' << std::endl; + } + + if (vals[i + 2] % 256 == 0) { + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2, GGML_TYPE_Q4_K); + std::cerr << '\n'; + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2, GGML_TYPE_Q4_K); + std::cerr << '\n'; + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1, GGML_TYPE_Q4_K); + ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2, GGML_TYPE_Q4_K); + std::cerr << '\n' << std::endl; + } + } + + GGML_ABORT("fatal error"); +#endif + + if (ctx->prealloc_x == nullptr || (ctx->prealloc_size_x > 0 && ctx->prealloc_x->size < ctx->prealloc_size_x)) { + VK_LOG_MEMORY("ggml_vk_preallocate_buffers(x_size: " << ctx->prealloc_size_x << ")"); + // Resize buffer + if (ctx->prealloc_x != nullptr) { + ggml_vk_destroy_buffer(ctx->prealloc_x); + } + ctx->prealloc_x = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_x); + } + if (ctx->prealloc_y == nullptr || (ctx->prealloc_size_y > 0 && ctx->prealloc_y->size < ctx->prealloc_size_y)) { + VK_LOG_MEMORY("ggml_vk_preallocate_buffers(y_size: " << ctx->prealloc_size_y << ")"); + // Resize buffer + if (ctx->prealloc_y != nullptr) { + ggml_vk_destroy_buffer(ctx->prealloc_y); + } + ctx->prealloc_y = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_y); + } + if (ctx->prealloc_split_k == nullptr || (ctx->prealloc_size_split_k > 0 && ctx->prealloc_split_k->size < ctx->prealloc_size_split_k)) { + VK_LOG_MEMORY("ggml_vk_preallocate_buffers(split_k_size: " << ctx->prealloc_size_split_k << ")"); + // Resize buffer + if (ctx->prealloc_split_k != nullptr) { + ggml_vk_destroy_buffer(ctx->prealloc_split_k); + } + ctx->prealloc_split_k = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_split_k); + } +} + +static bool ggml_vk_compute_forward(ggml_backend_vk_context* ctx, ggml_tensor* tensor, int tensor_idx, bool use_fence); + +// Returns true if node has enqueued work into the queue, false otherwise +// If submit is true the current all operations queued so far are being submitted to Vulkan to overlap cmdlist creation and GPU execution. +static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_tensor * node, int node_idx, ggml_tensor *node_begin, int node_idx_begin, bool dryrun, bool last_node, bool submit){ + if (ggml_is_empty(node) || !node->buffer) { + return false; + } + + VK_LOG_DEBUG("ggml_vk_build_graph(" << node << ", " << ggml_op_name(node->op) << ")"); + ctx->semaphore_idx = 0; + + const ggml_tensor * src0 = node->src[0]; + const ggml_tensor * src1 = node->src[1]; + const ggml_tensor * src2 = node->src[2]; + const ggml_tensor * src3 = node->src[3]; + + switch (node->op) { + // Return on empty ops to avoid generating a compute_ctx and setting exit_tensor + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_NONE: + return false; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(node)) { + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_GELU: + case GGML_UNARY_OP_GELU_QUICK: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_TANH: + break; + default: + return false; + } + break; + case GGML_OP_REPEAT: + case GGML_OP_GET_ROWS: + case GGML_OP_ADD: + case GGML_OP_ACC: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_DUP: + case GGML_OP_NORM: + case GGML_OP_GROUP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_SOFT_MAX: + case GGML_OP_ROPE: + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + case GGML_OP_ARGSORT: + case GGML_OP_SUM_ROWS: + case GGML_OP_IM2COL: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_POOL_2D: + case GGML_OP_RWKV_WKV6: + case GGML_OP_LEAKY_RELU: + case GGML_OP_FLASH_ATTN_EXT: + break; + default: + std::cerr << "ggml_vulkan: Error: Missing op: " << ggml_op_name(node->op) << std::endl; + GGML_ABORT("fatal error"); + return false; + } + + vk_context compute_ctx; + + if (!dryrun) { + if (ctx->compute_ctx.expired()) { + compute_ctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); + ctx->compute_ctx = compute_ctx; + ggml_vk_ctx_begin(ctx->device, compute_ctx); + } else { + compute_ctx = ctx->compute_ctx.lock(); + } + } else { + switch (node->op) { + case GGML_OP_REPEAT: + case GGML_OP_ACC: + case GGML_OP_GET_ROWS: + case GGML_OP_ADD: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_DUP: + case GGML_OP_NORM: + case GGML_OP_GROUP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_UNARY: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_SOFT_MAX: + case GGML_OP_ROPE: + case GGML_OP_ARGSORT: + case GGML_OP_SUM_ROWS: + case GGML_OP_IM2COL: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_POOL_2D: + case GGML_OP_LEAKY_RELU: + { + // These operations all go through ggml_vk_op_f32, so short-circuit and + // do the only thing needed for the dryrun. + vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, src2, node, node->op); + ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); + return false; + } + default: + break; + } + } + + switch (node->op) { + case GGML_OP_REPEAT: + ggml_vk_repeat(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_ACC: + ggml_vk_acc(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_GET_ROWS: + ggml_vk_get_rows(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_ADD: + ggml_vk_add(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_MUL: + ggml_vk_mul(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_DIV: + ggml_vk_div(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_CONCAT: + ggml_vk_concat(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_UPSCALE: + ggml_vk_upscale(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_SCALE: + ggml_vk_scale(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_SQR: + ggml_vk_sqr(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_SIN: + ggml_vk_sin(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_COS: + ggml_vk_cos(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_CLAMP: + ggml_vk_clamp(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_PAD: + ggml_vk_pad(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_DUP: + ggml_vk_cpy(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_NORM: + ggml_vk_norm(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_GROUP_NORM: + ggml_vk_group_norm(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_RMS_NORM: + ggml_vk_rms_norm(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(node)) { + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_GELU: + case GGML_UNARY_OP_GELU_QUICK: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_TANH: + ggml_vk_unary(ctx, compute_ctx, src0, node, dryrun); + break; + default: + return false; + } + break; + case GGML_OP_DIAG_MASK_INF: + ggml_vk_diag_mask_inf(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_SOFT_MAX: + ggml_vk_soft_max(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_ROPE: + ggml_vk_rope(ctx, compute_ctx, src0, src1, src2, node, dryrun); + + break; + case GGML_OP_ARGSORT: + ggml_vk_argsort(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_SUM_ROWS: + ggml_vk_sum_rows(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_IM2COL: + ggml_vk_im2col(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_TIMESTEP_EMBEDDING: + ggml_vk_timestep_embedding(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_POOL_2D: + ggml_vk_pool_2d(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_LEAKY_RELU: + ggml_vk_leaky_relu(ctx, compute_ctx, src0, node, dryrun); + + break; + case GGML_OP_MUL_MAT: + ggml_vk_mul_mat(ctx, compute_ctx, src0, src1, node, dryrun); + + break; + case GGML_OP_MUL_MAT_ID: + ggml_vk_mul_mat_id(ctx, compute_ctx, src0, src1, src2, node, dryrun); + + break; + + case GGML_OP_FLASH_ATTN_EXT: + ggml_vk_flash_attn(ctx, compute_ctx, src0, src1, src2, src3, node, dryrun); + + break; + + case GGML_OP_RWKV_WKV6: + ggml_vk_rwkv_wkv6(ctx, compute_ctx, node, dryrun); + + break; + default: + return false; + } + + if (dryrun) { + return false; + } + + ctx->tensor_ctxs[node_idx] = compute_ctx; + +#if defined(GGML_VULKAN_CHECK_RESULTS) || defined(GGML_VULKAN_PERF) + // Force context reset on each node so that each tensor ends up in its own context + // and can be run and compared to its CPU equivalent separately + last_node = true; +#endif + + if (submit || last_node) { + ggml_vk_ctx_end(compute_ctx); + + // TODO probably it'd be better to pass a exit_node flag to ggml_vk_compute_forward + if (last_node) { + compute_ctx->exit_tensor_idx = node_idx_begin; + } + else { + compute_ctx->exit_tensor_idx = -1; + } + + ctx->compute_ctx.reset(); + + bool ok = ggml_vk_compute_forward(ctx, node_begin, node_idx_begin, false); + if (!ok) { + if (node->op == GGML_OP_UNARY) { + std::cerr << __func__ << ": error: op not supported UNARY " << node->name << " (" << ggml_unary_op_name(static_cast(node->op_params[0])) << ")" << std::endl; + } + else { + std::cerr << __func__ << ": error: op not supported " << node->name << " (" << ggml_op_name(node->op) << ")" << std::endl; + } + } + + } + return true; +} + +static bool ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_tensor * tensor, int tensor_idx, bool use_fence = true){ + ggml_backend_buffer * buf = nullptr; + + switch (tensor->op) { + case GGML_OP_ADD: + case GGML_OP_ACC: + case GGML_OP_GET_ROWS: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_DUP: + case GGML_OP_NORM: + case GGML_OP_GROUP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_SOFT_MAX: + case GGML_OP_ROPE: + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_NONE: + case GGML_OP_ARGSORT: + case GGML_OP_SUM_ROWS: + case GGML_OP_IM2COL: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_POOL_2D: + case GGML_OP_RWKV_WKV6: + case GGML_OP_LEAKY_RELU: + case GGML_OP_REPEAT: + buf = tensor->buffer; + + break; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(tensor)) { + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_GELU: + case GGML_UNARY_OP_GELU_QUICK: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_TANH: + buf = tensor->buffer; + break; + default: + return false; + } + break; + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + case GGML_OP_FLASH_ATTN_EXT: + buf = tensor->buffer; + + break; + default: + return false; + } + + if (buf == nullptr) { + return false; + } + + VK_LOG_DEBUG("ggml_vk_compute_forward(" << tensor << ", name=" << tensor->name << ", op=" << ggml_op_name(tensor->op) << ", type=" << tensor->type << ", ne0=" << tensor->ne[0] << ", ne1=" << tensor->ne[1] << ", ne2=" << tensor->ne[2] << ", ne3=" << tensor->ne[3] << ", nb0=" << tensor->nb[0] << ", nb1=" << tensor->nb[1] << ", nb2=" << tensor->nb[2] << ", nb3=" << tensor->nb[3] << ", view_src=" << tensor->view_src << ", view_offs=" << tensor->view_offs << ")"); + + vk_context subctx = ctx->tensor_ctxs[tensor_idx].lock(); + + // always wait for the GPU work to be done for the last submit + if (tensor_idx == subctx->exit_tensor_idx) { + use_fence = true; + } + + // Only run if ctx hasn't been submitted yet + if (!subctx->seqs.empty()) { +#ifdef GGML_VULKAN_CHECK_RESULTS + ggml_vk_check_results_0(tensor); + use_fence = true; +#endif + + // Do staging buffer copies + for (auto& cpy : subctx->in_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + + ggml_vk_submit(subctx, use_fence ? ctx->fence : vk::Fence{}); + + if (use_fence) { + VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_compute_forward waitForFences"); + + ctx->device->device.resetFences({ ctx->fence }); + } +#ifdef GGML_VULKAN_CHECK_RESULTS + ggml_vk_check_results_1(tensor); +#endif + } + + if (tensor_idx == subctx->exit_tensor_idx) { + // Do staging buffer copies + for (auto& cpy : subctx->out_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + subctx->in_memcpys.clear(); + subctx->out_memcpys.clear(); + } + + return true; +} + +// Clean up after graph processing is done +static void ggml_vk_graph_cleanup(ggml_backend_vk_context * ctx) { + VK_LOG_DEBUG("ggml_vk_graph_cleanup()"); + for (auto& buffer : ctx->gc.temp_buffers) { + ggml_vk_pool_free(ctx, buffer); + } + ctx->gc.temp_buffers.clear(); + + for (auto& dsr : ctx->device->pipeline_descriptor_set_requirements) { + vk_pipeline_ref plr = ctx->device->pipelines[dsr.first]; + + if (plr.expired()) { + continue; + } + + vk_pipeline pl = plr.lock(); + ggml_pipeline_cleanup(pl); + } + + ggml_vk_queue_cleanup(ctx->device, ctx->device->compute_queue); + ggml_vk_queue_cleanup(ctx->device, ctx->device->transfer_queue); + + for (size_t i = 0; i < ctx->gc.semaphores.size(); i++) { + ctx->device->device.destroySemaphore({ ctx->gc.semaphores[i].s }); + } + ctx->gc.semaphores.clear(); + + for (size_t i = 0; i < ctx->gc.tl_semaphores.size(); i++) { + ctx->device->device.destroySemaphore({ ctx->gc.tl_semaphores[i].s }); + } + ctx->gc.tl_semaphores.clear(); + ctx->semaphore_idx = 0; + + ctx->event_idx = 0; + + for (auto& event : ctx->gc.events) { + ctx->device->device.resetEvent(event); + } + + ctx->tensor_ctxs.clear(); + ctx->gc.contexts.clear(); + ctx->device->pipeline_descriptor_set_requirements.clear(); +} + +// Clean up on backend free +static void ggml_vk_cleanup(ggml_backend_vk_context * ctx) { + VK_LOG_DEBUG("ggml_vk_cleanup(" << ctx->name << ")"); + ggml_vk_graph_cleanup(ctx); + + ggml_vk_destroy_buffer(ctx->prealloc_x); + ggml_vk_destroy_buffer(ctx->prealloc_y); + ggml_vk_destroy_buffer(ctx->prealloc_split_k); + + for (auto& buffer : ctx->buffer_pool) { + ggml_vk_destroy_buffer(buffer); + } + + ctx->prealloc_size_x = 0; + ctx->prealloc_size_y = 0; + ctx->prealloc_size_split_k = 0; + + for (auto& event : ctx->gc.events) { + ctx->device->device.destroyEvent(event); + } + ctx->gc.events.clear(); + + ctx->device->device.destroyFence(ctx->fence); +} + +static int ggml_vk_get_device_count() { + ggml_vk_instance_init(); + + return vk_instance.device_indices.size(); +} + +static void ggml_vk_get_device_description(int device, char * description, size_t description_size) { + ggml_vk_instance_init(); + + std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); + + vk::PhysicalDeviceProperties props; + devices[device].getProperties(&props); + + snprintf(description, description_size, "%s", props.deviceName.data()); +} + +// backend interface + +#define UNUSED GGML_UNUSED + +// device backend + +static bool ggml_backend_buffer_is_vk(ggml_backend_buffer_t buffer) { + return buffer->buft->iface.get_name == ggml_backend_vk_buffer_type_name; +} + +static void ggml_backend_vk_buffer_free_buffer(ggml_backend_buffer_t buffer) { + VK_LOG_MEMORY("ggml_backend_vk_buffer_free_buffer()"); + ggml_backend_vk_buffer_context * ctx = (ggml_backend_vk_buffer_context *)buffer->context; + ggml_vk_destroy_buffer(ctx->dev_buffer); + delete ctx; +} + +static void * ggml_backend_vk_buffer_get_base(ggml_backend_buffer_t buffer) { + return vk_ptr_base; + + UNUSED(buffer); +} + +static void ggml_backend_vk_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + VK_LOG_DEBUG("ggml_backend_vk_buffer_init_tensor(" << buffer << " (" << buffer->context << "), " << tensor << ")"); + if (tensor->view_src != nullptr) { + GGML_ASSERT(tensor->view_src->buffer->buft == buffer->buft); + } +} + +static void ggml_backend_vk_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + VK_LOG_DEBUG("ggml_backend_vk_buffer_set_tensor(" << buffer << ", " << tensor << ", " << data << ", " << offset << ", " << size << ")"); + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)buffer->context; + vk_buffer buf = buf_ctx->dev_buffer; + + ggml_vk_buffer_write(buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); +} + +static void ggml_backend_vk_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + VK_LOG_DEBUG("ggml_backend_vk_buffer_get_tensor(" << buffer << ", " << tensor << ", " << data << ", " << offset << ", " << size << ")"); + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)buffer->context; + + vk_buffer buf = buf_ctx->dev_buffer; + + ggml_vk_buffer_read(buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); +} + +static bool ggml_backend_vk_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) { + if (ggml_backend_buffer_is_vk(src->buffer)) { + ggml_backend_vk_buffer_context * src_buf_ctx = (ggml_backend_vk_buffer_context *)src->buffer->context; + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + + vk_buffer src_buf = src_buf_ctx->dev_buffer; + vk_buffer dst_buf = dst_buf_ctx->dev_buffer; + + ggml_vk_buffer_copy(dst_buf, vk_tensor_offset(dst) + dst->view_offs, src_buf, vk_tensor_offset(src) + src->view_offs, ggml_nbytes(src)); + + return true; + } + return false; + + UNUSED(buffer); +} + +static void ggml_backend_vk_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + ggml_backend_vk_buffer_context * ctx = (ggml_backend_vk_buffer_context *)buffer->context; + + ggml_vk_buffer_memset(ctx->dev_buffer, 0, value, buffer->size); +} + +static ggml_backend_buffer_i ggml_backend_vk_buffer_interface = { + /* .free_buffer = */ ggml_backend_vk_buffer_free_buffer, + /* .get_base = */ ggml_backend_vk_buffer_get_base, + /* .init_tensor = */ ggml_backend_vk_buffer_init_tensor, + /* .memset_tensor = */ NULL, + /* .set_tensor = */ ggml_backend_vk_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_vk_buffer_get_tensor, + /* .cpy_tensor = */ ggml_backend_vk_buffer_cpy_tensor, + /* .clear = */ ggml_backend_vk_buffer_clear, + /* .reset = */ NULL, +}; + +// vk buffer type +static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft) { + ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *)buft->context; + + return ctx->name.c_str(); +} + +static ggml_backend_buffer_t ggml_backend_vk_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + VK_LOG_MEMORY("ggml_backend_vk_buffer_type_alloc_buffer(" << size << ")"); + ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; + + vk_buffer dev_buffer = nullptr; + try { + dev_buffer = ggml_vk_create_buffer_device(ctx->device, size); + } catch (const vk::SystemError& e) { + return nullptr; + } + + ggml_backend_vk_buffer_context * bufctx = new ggml_backend_vk_buffer_context(ctx->device, std::move(dev_buffer), ctx->name); + + return ggml_backend_buffer_init(buft, ggml_backend_vk_buffer_interface, bufctx, size); +} + +static size_t ggml_backend_vk_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; + return ctx->device->properties.limits.minStorageBufferOffsetAlignment; +} + +static size_t ggml_backend_vk_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { + ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; + return ctx->device->max_memory_allocation_size; +} + +static size_t ggml_backend_vk_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + return ggml_nbytes(tensor); + + UNUSED(buft); +} + +ggml_backend_buffer_type_t ggml_backend_vk_buffer_type(size_t dev_num) { + ggml_vk_instance_init(); + + VK_LOG_DEBUG("ggml_backend_vk_buffer_type(" << dev_num << ")"); + + vk_device dev = ggml_vk_get_device(dev_num); + + return &dev->buffer_type; +} + +// host buffer type + +static const char * ggml_backend_vk_host_buffer_type_name(ggml_backend_buffer_type_t buft) { + return GGML_VK_NAME "_Host"; + + UNUSED(buft); +} + +static const char * ggml_backend_vk_host_buffer_name(ggml_backend_buffer_t buffer) { + return GGML_VK_NAME "_Host"; + + UNUSED(buffer); +} + +static void ggml_backend_vk_host_buffer_free_buffer(ggml_backend_buffer_t buffer) { + VK_LOG_MEMORY("ggml_backend_vk_host_buffer_free_buffer()"); + ggml_vk_host_free(vk_instance.devices[0], buffer->context); +} + +static ggml_backend_buffer_t ggml_backend_vk_host_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + VK_LOG_MEMORY("ggml_backend_vk_host_buffer_type_alloc_buffer(" << size << ")"); + + size += 32; // Behave like the CPU buffer type + void * ptr = nullptr; + try { + ptr = ggml_vk_host_malloc(vk_instance.devices[0], size); + } catch (vk::SystemError& e) { + std::cerr << "ggml_vulkan: Failed to allocate pinned memory." << std::endl; + std::cerr << "ggml_vulkan: " << e.what() << std::endl; + // fallback to cpu buffer + return ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); + } + + ggml_backend_buffer_t buffer = ggml_backend_cpu_buffer_from_ptr(ptr, size); + buffer->buft = buft; + buffer->iface.free_buffer = ggml_backend_vk_host_buffer_free_buffer; + + return buffer; + + UNUSED(buft); +} + +static size_t ggml_backend_vk_host_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + return vk_instance.devices[0]->properties.limits.minMemoryMapAlignment; + + UNUSED(buft); +} + +// Should be changed to return device-specific host buffer type +// but that probably requires changes in llama.cpp +ggml_backend_buffer_type_t ggml_backend_vk_host_buffer_type() { + static struct ggml_backend_buffer_type ggml_backend_vk_buffer_type_host = { + /* .iface = */ { + /* .get_name = */ ggml_backend_vk_host_buffer_type_name, + /* .alloc_buffer = */ ggml_backend_vk_host_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_vk_host_buffer_type_get_alignment, + /* .get_max_size = */ NULL, // defaults to SIZE_MAX + /* .get_alloc_size = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size, + /* .is_host = */ ggml_backend_cpu_buffer_type()->iface.is_host, + }, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), 0), + /* .context = */ nullptr, + }; + + // Make sure device 0 is initialized + ggml_vk_instance_init(); + ggml_vk_get_device(0); + + return &ggml_backend_vk_buffer_type_host; +} + + +// backend + +static const char * ggml_backend_vk_name(ggml_backend_t backend) { + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + + return ctx->name.c_str(); +} + +static void ggml_backend_vk_free(ggml_backend_t backend) { + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + VK_LOG_DEBUG("ggml_backend_vk_free(" << ctx->name << ")"); + + ggml_vk_cleanup(ctx); + + delete ctx; + delete backend; +} + +static ggml_backend_buffer_type_t ggml_backend_vk_get_default_buffer_type(ggml_backend_t backend) { + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + + return &ctx->device->buffer_type; +} + +static void ggml_backend_vk_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + VK_LOG_DEBUG("ggml_backend_vk_set_tensor_async(" << size << ")"); + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + GGML_ASSERT((tensor->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || tensor->buffer->buft == ggml_backend_vk_host_buffer_type()) && "unsupported buffer type"); + + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; + + vk_context transfer_ctx; + + if (ctx->transfer_ctx.expired()) { + // Initialize new transfer context + transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); + ctx->transfer_ctx = transfer_ctx; + ggml_vk_ctx_begin(ctx->device, transfer_ctx); + } else { + transfer_ctx = ctx->transfer_ctx.lock(); + } + + vk_buffer buf = buf_ctx->dev_buffer; + + ggml_vk_buffer_write_async(transfer_ctx, buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); +} + +static void ggml_backend_vk_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + VK_LOG_DEBUG("ggml_backend_vk_get_tensor_async(" << size << ")"); + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + GGML_ASSERT((tensor->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || tensor->buffer->buft == ggml_backend_vk_host_buffer_type()) && "unsupported buffer type"); + + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; + + vk_context transfer_ctx; + + if (ctx->transfer_ctx.expired()) { + // Initialize new transfer context + transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); + ctx->transfer_ctx = transfer_ctx; + ggml_vk_ctx_begin(ctx->device, transfer_ctx); + } else { + transfer_ctx = ctx->transfer_ctx.lock(); + } + + vk_buffer buf = buf_ctx->dev_buffer; + + ggml_vk_buffer_read_async(transfer_ctx, buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); +} + +static bool ggml_backend_vk_cpy_tensor_async(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst) { + VK_LOG_DEBUG("ggml_backend_vk_cpy_tensor_async()"); + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + if ((dst->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || dst->buffer->buft == ggml_backend_vk_host_buffer_type()) && ggml_backend_buffer_is_vk(src->buffer)) { + ggml_backend_vk_buffer_context * src_buf_ctx = (ggml_backend_vk_buffer_context *)src->buffer->context; + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + + vk_context transfer_ctx; + + if (ctx->transfer_ctx.expired()) { + // Initialize new transfer context + transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); + ctx->transfer_ctx = transfer_ctx; + ggml_vk_ctx_begin(ctx->device, transfer_ctx); + } else { + transfer_ctx = ctx->transfer_ctx.lock(); + } + + vk_buffer src_buf = src_buf_ctx->dev_buffer; + vk_buffer dst_buf = dst_buf_ctx->dev_buffer; + + ggml_vk_buffer_copy_async(transfer_ctx, dst_buf, vk_tensor_offset(dst) + dst->view_offs, src_buf, vk_tensor_offset(src) + src->view_offs, ggml_nbytes(src)); + return true; + } + + return false; +} + +static void ggml_backend_vk_synchronize(ggml_backend_t backend) { + VK_LOG_DEBUG("ggml_backend_vk_synchronize()"); + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + if(ctx->transfer_ctx.expired()) { + return; + } + + vk_context transfer_ctx = ctx->transfer_ctx.lock(); + + ggml_vk_ctx_end(transfer_ctx); + + for (auto& cpy : transfer_ctx->in_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + + ggml_vk_submit(transfer_ctx, ctx->fence); + VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_backend_vk_synchronize waitForFences"); + ctx->device->device.resetFences({ ctx->fence }); + + for (auto& cpy : transfer_ctx->out_memcpys) { + memcpy(cpy.dst, cpy.src, cpy.n); + } + + ctx->transfer_ctx.reset(); +} + +static bool ggml_vk_is_empty(ggml_tensor * node) { + return ggml_is_empty(node) || node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE; +} + +static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { + VK_LOG_DEBUG("ggml_backend_vk_graph_compute(" << cgraph->n_nodes << " nodes)"); + ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; + + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_vk_build_graph(ctx, cgraph->nodes[i], i, nullptr, 0, true, false, false); + } + ggml_vk_preallocate_buffers(ctx); + ggml_pipeline_allocate_descriptor_sets(ctx->device); + + int last_node = cgraph->n_nodes - 1; + + // If the last op in the cgraph isn't backend GPU, the command buffer doesn't get closed properly + while (last_node > 0 && ggml_vk_is_empty(cgraph->nodes[last_node])) { + last_node -= 1; + } + + // Reserve tensor context space for all nodes + ctx->tensor_ctxs.resize(cgraph->n_nodes); + + bool first_node_in_batch = true; // true if next node will be first node in a batch + int submit_node_idx = 0; // index to first node in a batch + + // Submit work every nodes_per_submit nodes to overlap CPU cmdbuffer generation with GPU execution. + // Start with a smaller count to get work submitted right away, and increase it after each submit. + int nodes_per_submit = 20; + int submitted_nodes = 0; + int submit_count = 0; + for (int i = 0; i < cgraph->n_nodes; i++) { + if (first_node_in_batch) { + submit_node_idx = i; + } + + bool submit = (submitted_nodes >= nodes_per_submit) || (i == last_node); + + bool enqueued = ggml_vk_build_graph(ctx, cgraph->nodes[i], i, cgraph->nodes[submit_node_idx], submit_node_idx, false, i == last_node, submit); + + if (enqueued) { + ++submitted_nodes; + +#ifndef GGML_VULKAN_CHECK_RESULTS + if (first_node_in_batch) { + first_node_in_batch = false; + } +#endif + } + + if (submit) { + first_node_in_batch = true; + submitted_nodes = 0; + switch (submit_count) { + case 0: + nodes_per_submit = 50; + break; + default: + nodes_per_submit = 100; + break; + } + submit_count++; + } + } + +#ifdef GGML_VULKAN_PERF + ctx->device->perf_logger->print_timings(); +#endif + + ggml_vk_graph_cleanup(ctx); + + return GGML_STATUS_SUCCESS; + + UNUSED(backend); +} + +// TODO: enable async and synchronize +static ggml_backend_i ggml_backend_vk_interface = { + /* .get_name = */ ggml_backend_vk_name, + /* .free = */ ggml_backend_vk_free, + /* .set_tensor_async = */ NULL, // ggml_backend_vk_set_tensor_async, + /* .get_tensor_async = */ NULL, // ggml_backend_vk_get_tensor_async, + /* .cpy_tensor_async = */ NULL, // ggml_backend_vk_cpy_tensor_async, + /* .synchronize = */ NULL, // ggml_backend_vk_synchronize, + /* .graph_plan_create = */ NULL, + /* .graph_plan_free = */ NULL, + /* .graph_plan_update = */ NULL, + /* .graph_plan_compute = */ NULL, + /* .graph_compute = */ ggml_backend_vk_graph_compute, + /* .event_record = */ NULL, + /* .event_wait = */ NULL, +}; + +static ggml_guid_t ggml_backend_vk_guid() { + static ggml_guid guid = { 0xb8, 0xf7, 0x4f, 0x86, 0x40, 0x3c, 0xe1, 0x02, 0x91, 0xc8, 0xdd, 0xe9, 0x02, 0x3f, 0xc0, 0x2b }; + return &guid; +} + +ggml_backend_t ggml_backend_vk_init(size_t dev_num) { + VK_LOG_DEBUG("ggml_backend_vk_init(" << dev_num << ")"); + + ggml_backend_vk_context * ctx = new ggml_backend_vk_context; + ggml_vk_init(ctx, dev_num); + + ggml_backend_t vk_backend = new ggml_backend { + /* .guid = */ ggml_backend_vk_guid(), + /* .interface = */ ggml_backend_vk_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), dev_num), + /* .context = */ ctx, + }; + + return vk_backend; +} + +bool ggml_backend_is_vk(ggml_backend_t backend) { + return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_vk_guid()); +} + +int ggml_backend_vk_get_device_count() { + return ggml_vk_get_device_count(); +} + +void ggml_backend_vk_get_device_description(int device, char * description, size_t description_size) { + GGML_ASSERT(device < (int) vk_instance.device_indices.size()); + int dev_idx = vk_instance.device_indices[device]; + ggml_vk_get_device_description(dev_idx, description, description_size); +} + +void ggml_backend_vk_get_device_memory(int device, size_t * free, size_t * total) { + GGML_ASSERT(device < (int) vk_instance.device_indices.size()); + + vk::PhysicalDevice vkdev = vk_instance.instance.enumeratePhysicalDevices()[vk_instance.device_indices[device]]; + + vk::PhysicalDeviceMemoryProperties memprops = vkdev.getMemoryProperties(); + + for (const vk::MemoryHeap& heap : memprops.memoryHeaps) { + if (heap.flags & vk::MemoryHeapFlagBits::eDeviceLocal) { + *total = heap.size; + *free = heap.size; + break; + } + } +} + +////////////////////////// + +struct ggml_backend_vk_device_context { + size_t device; + std::string name; + std::string description; +}; + +static const char * ggml_backend_vk_device_get_name(ggml_backend_dev_t dev) { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + return ctx->name.c_str(); +} + +static const char * ggml_backend_vk_device_get_description(ggml_backend_dev_t dev) { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + return ctx->description.c_str(); +} + +static void ggml_backend_vk_device_get_memory(ggml_backend_dev_t device, size_t * free, size_t * total) { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)device->context; + ggml_backend_vk_get_device_memory(ctx->device, free, total); +} + +static ggml_backend_buffer_type_t ggml_backend_vk_device_get_buffer_type(ggml_backend_dev_t dev) { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + return ggml_backend_vk_buffer_type(ctx->device); +} + +static ggml_backend_buffer_type_t ggml_backend_vk_device_get_host_buffer_type(ggml_backend_dev_t dev) { + UNUSED(dev); + return ggml_backend_vk_host_buffer_type(); +} + +static enum ggml_backend_dev_type ggml_backend_vk_device_get_type(ggml_backend_dev_t dev) { + UNUSED(dev); + return GGML_BACKEND_DEVICE_TYPE_GPU; +} + +static void ggml_backend_vk_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props) { + props->name = ggml_backend_vk_device_get_name(dev); + props->description = ggml_backend_vk_device_get_description(dev); + props->type = ggml_backend_vk_device_get_type(dev); + ggml_backend_vk_device_get_memory(dev, &props->memory_free, &props->memory_total); + props->caps = { + /* .async = */ false, + /* .host_buffer = */ true, + /* .buffer_from_host_ptr = */ false, + /* .events = */ false, + }; +} + +static ggml_backend_t ggml_backend_vk_device_init(ggml_backend_dev_t dev, const char * params) { + UNUSED(params); + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + return ggml_backend_vk_init(ctx->device); +} + +static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + switch (op->op) { + case GGML_OP_UNARY: + switch (ggml_get_unary_op(op)) { + case GGML_UNARY_OP_GELU: + case GGML_UNARY_OP_GELU_QUICK: + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_TANH: + return ggml_is_contiguous(op->src[0]); + default: + return false; + } + break; + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + const vk_device& device = ggml_vk_get_device(ctx->device); + if (op->op == GGML_OP_MUL_MAT_ID && !device->mul_mat_id_s && !device->mul_mat_id_m && !device->mul_mat_id_l) { + // If there's not enough shared memory for row_ids and the result tile, fallback to CPU + return false; + } + switch (op->src[0]->type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return false; + } + struct ggml_tensor * a; + struct ggml_tensor * b; + if (op->op == GGML_OP_MUL_MAT) { + a = op->src[0]; + b = op->src[1]; + } else { + a = op->src[2]; + b = op->src[1]; + } + if (a->ne[3] != b->ne[3]) { + return false; + } + if (!(ggml_vk_dim01_contiguous(op->src[0]) || op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) || + !(ggml_vk_dim01_contiguous(op->src[1]) || op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16)) { + return false; + } + + return true; + } break; + case GGML_OP_FLASH_ATTN_EXT: + { + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + if (!ggml_vk_get_device(ctx->device)->coopmat2) { + return false; + } + switch (op->src[0]->ne[0]) { + case 64: + case 80: + case 96: + case 112: + case 128: + case 256: + break; + default: + return false; + } + if (op->src[0]->type != GGML_TYPE_F32) { + return false; + } + if (op->type != GGML_TYPE_F32) { + return false; + } + if (op->src[3] && op->src[3]->type != GGML_TYPE_F16) { + return false; + } + // It's straightforward to support different K/V dequant, but would + // significantly increase the number of pipelines + if (op->src[1]->type != op->src[2]->type) { + return false; + } + switch (op->src[1]->type) { + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + // K dequants currently disabled because D dimension is rounded up to 256 and runs inefficiently + //case GGML_TYPE_Q2_K: + //case GGML_TYPE_Q3_K: + //case GGML_TYPE_Q4_K: + //case GGML_TYPE_Q5_K: + //case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + break; + default: + return false; + } + return true; + } + case GGML_OP_GET_ROWS: + { + switch (op->src[0]->type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_IQ4_NL: + return true; + default: + return false; + } + } break; + case GGML_OP_CONT: + case GGML_OP_CPY: + case GGML_OP_DUP: + { + ggml_type src0_type = op->src[0]->type; + ggml_type src1_type = op->src[1] != nullptr ? op->src[1]->type : src0_type; + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F16) { + return true; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return true; + } + return false; + } break; + case GGML_OP_REPEAT: + return ggml_type_size(op->type) == sizeof(float) && ggml_type_size(op->src[0]->type) == sizeof(float); + case GGML_OP_ROPE: + { + const int mode = ((const int32_t *) op->op_params)[2]; + if (mode & GGML_ROPE_TYPE_MROPE) { + return false; + } + if (mode & GGML_ROPE_TYPE_VISION) { + return false; + } + return ggml_is_contiguous(op->src[0]); + } + case GGML_OP_NONE: + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_NORM: + case GGML_OP_GROUP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_ADD: + case GGML_OP_ACC: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_CONCAT: + case GGML_OP_UPSCALE: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_PAD: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_SOFT_MAX: + case GGML_OP_ARGSORT: + case GGML_OP_SUM_ROWS: + case GGML_OP_IM2COL: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_POOL_2D: + case GGML_OP_RWKV_WKV6: + case GGML_OP_LEAKY_RELU: + return true; + default: + return false; + } + + UNUSED(dev); +} + +static bool ggml_backend_vk_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { + if (buft->iface.get_name != ggml_backend_vk_buffer_type_name) { + return false; + } + + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; + ggml_backend_vk_buffer_type_context * buft_ctx = (ggml_backend_vk_buffer_type_context *)buft->context; + + return buft_ctx->device->idx == ctx->device; +} + +static bool ggml_backend_vk_device_offload_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + const int min_batch_size = 32; + + return (op->ne[1] >= min_batch_size && op->op != GGML_OP_GET_ROWS) || + (op->ne[2] >= min_batch_size && op->op == GGML_OP_MUL_MAT_ID); + + UNUSED(dev); +} + +static const struct ggml_backend_device_i ggml_backend_vk_device_i = { + /* .get_name = */ ggml_backend_vk_device_get_name, + /* .get_description = */ ggml_backend_vk_device_get_description, + /* .get_memory = */ ggml_backend_vk_device_get_memory, + /* .get_type = */ ggml_backend_vk_device_get_type, + /* .get_props = */ ggml_backend_vk_device_get_props, + /* .init_backend = */ ggml_backend_vk_device_init, + /* .get_buffer_type = */ ggml_backend_vk_device_get_buffer_type, + /* .get_host_buffer_type = */ ggml_backend_vk_device_get_host_buffer_type, + /* .buffer_from_host_ptr = */ NULL, + /* .supports_op = */ ggml_backend_vk_device_supports_op, + /* .supports_buft = */ ggml_backend_vk_device_supports_buft, + /* .offload_op = */ ggml_backend_vk_device_offload_op, + /* .event_new = */ NULL, + /* .event_free = */ NULL, + /* .event_synchronize = */ NULL, +}; + +static const char * ggml_backend_vk_reg_get_name(ggml_backend_reg_t reg) { + UNUSED(reg); + return GGML_VK_NAME; +} + +static size_t ggml_backend_vk_reg_get_device_count(ggml_backend_reg_t reg) { + UNUSED(reg); + return ggml_backend_vk_get_device_count(); +} + +static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg, size_t device) { + static std::vector devices; + + static bool initialized = false; + + { + static std::mutex mutex; + std::lock_guard lock(mutex); + if (!initialized) { + for (int i = 0; i < ggml_backend_vk_get_device_count(); i++) { + ggml_backend_vk_device_context * ctx = new ggml_backend_vk_device_context; + char desc[256]; + ggml_backend_vk_get_device_description(i, desc, sizeof(desc)); + ctx->device = i; + ctx->name = GGML_VK_NAME + std::to_string(i); + ctx->description = desc; + devices.push_back(new ggml_backend_device { + /* .iface = */ ggml_backend_vk_device_i, + /* .reg = */ reg, + /* .context = */ ctx, + }); + } + initialized = true; + } + } + + GGML_ASSERT(device < devices.size()); + return devices[device]; +} + +static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = { + /* .get_name = */ ggml_backend_vk_reg_get_name, + /* .get_device_count = */ ggml_backend_vk_reg_get_device_count, + /* .get_device = */ ggml_backend_vk_reg_get_device, + /* .get_proc_address = */ NULL, +}; + +ggml_backend_reg_t ggml_backend_vk_reg() { + static ggml_backend_reg reg = { + /* .api_version = */ GGML_BACKEND_API_VERSION, + /* .iface = */ ggml_backend_vk_reg_i, + /* .context = */ nullptr, + }; + + return ® +} + +// Extension availability +static bool ggml_vk_instance_validation_ext_available(const std::vector& instance_extensions) { +#ifdef GGML_VULKAN_VALIDATE + bool portability_enumeration_ext = false; + // Check for portability enumeration extension for MoltenVK support + for (const auto& properties : instance_extensions) { + if (strcmp("VK_KHR_portability_enumeration", properties.extensionName) == 0) { + return true; + } + } + if (!portability_enumeration_ext) { + std::cerr << "ggml_vulkan: WARNING: Instance extension VK_KHR_portability_enumeration not found." << std::endl; + } +#endif + return false; + + UNUSED(instance_extensions); +} +static bool ggml_vk_instance_portability_enumeration_ext_available(const std::vector& instance_extensions) { +#ifdef __APPLE__ + bool portability_enumeration_ext = false; + // Check for portability enumeration extension for MoltenVK support + for (const auto& properties : instance_extensions) { + if (strcmp("VK_KHR_portability_enumeration", properties.extensionName) == 0) { + return true; + } + } + if (!portability_enumeration_ext) { + std::cerr << "ggml_vulkan: WARNING: Instance extension VK_KHR_portability_enumeration not found." << std::endl; + } +#endif + return false; + + UNUSED(instance_extensions); +} + +static bool ggml_vk_khr_cooperative_matrix_support(const vk::PhysicalDeviceProperties& props, const vk::PhysicalDeviceDriverProperties& driver_props) { + switch (props.vendorID) { + case VK_VENDOR_ID_INTEL: + // Intel drivers don't support coopmat properly yet + return false; + case VK_VENDOR_ID_AMD: + if (driver_props.driverID == vk::DriverId::eAmdProprietary || driver_props.driverID == vk::DriverId::eAmdOpenSource) { + // Workaround for AMD proprietary driver reporting support on all GPUs + const std::string name = props.deviceName; + return name.rfind("AMD Radeon RX 7", 0) == 0 || name.rfind("AMD Radeon(TM) RX 7", 0) == 0 || // RDNA 3 consumer GPUs + name.rfind("AMD Radeon PRO W7", 0) == 0 || name.rfind("AMD Radeon(TM) PRO W7", 0) == 0 || // RDNA 3 workstation GPUs + name.rfind("AMD Radeon 7", 0) == 0 || name.rfind("AMD Radeon(TM) 7", 0) == 0; // RDNA 3 APUs + } + return true; + default: + return true; + } +} + +// checks + +#ifdef GGML_VULKAN_CHECK_RESULTS +static void ggml_vk_print_graph_origin(const ggml_tensor * tensor, std::vector& done, int level = 0) { + if (std::find(done.begin(), done.end(), tensor) != done.end() || level > 10) { + return; + } + for (int j = 0; j < level; j++) { + std::cerr << " "; + } + std::cerr << ggml_op_name(tensor->op) << " gpu=" << (tensor->extra != nullptr) << std::endl; + + done.push_back(tensor); + + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (tensor->src[i] != nullptr) { + ggml_vk_print_graph_origin(tensor->src[i], done, level + 1); + } + } +} + +static void ggml_vk_print_tensor_area(const ggml_tensor * tensor, const void * data, int i0, int i1, int i2, int i3) { + if (tensor->type != GGML_TYPE_F32 && tensor->type != GGML_TYPE_F16 && tensor->type != GGML_TYPE_I32) { + return; + } + i0 = std::max(i0, 5); + i1 = std::max(i1, 5); + i2 = std::max(i2, 0); + i3 = std::max(i3, 0); + fprintf(stderr, " "); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + fprintf(stderr, "%7d ", idx1); + } + fprintf(stderr, "\n"); + for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { + fprintf(stderr, "%7d: ", idx0); + for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { + if (idx0 >= 0 && idx0 < tensor->ne[0] && idx1 >= 0 && idx1 < tensor->ne[1] && i2 >= 0 && i2 < tensor->ne[2] && i3 >= 0 && i3 < tensor->ne[3]) { + float val; + if (tensor->type == GGML_TYPE_F32) { + val = *(const float *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); + } else if (tensor->type == GGML_TYPE_F16) { + val = ggml_fp16_to_fp32(*(const ggml_fp16_t *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0])); + } else if (tensor->type == GGML_TYPE_I32) { + val = *(const int32_t *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); + } else { + GGML_ABORT("fatal error"); + } + fprintf(stderr, "% 7.2f ", val); + } else { + fprintf(stderr, " "); + } + } + fprintf(stderr, "\n"); + } +} + +static void ggml_vk_print_tensor(const ggml_tensor * tensor, const char * name) { + void * tensor_data = tensor->data; + + const bool is_gpu = tensor->buffer != nullptr && ggml_backend_buffer_is_vk(tensor->buffer); + + if (is_gpu) { + const size_t tensor_size = ggml_nbytes(tensor); + tensor_data = malloc(tensor_size); + + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; + + vk_buffer buffer_gpu = buf_ctx->dev_buffer; + ggml_vk_buffer_read(buffer_gpu, vk_tensor_offset(tensor) + tensor->view_offs, tensor_data, tensor_size); + } + + std::cerr << "TENSOR CHECK " << name << " (" << tensor->name << "): " << ggml_op_name(tensor->op) << std::endl; + std::cerr << "tensor=" << tensor << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << std::endl; + if (tensor->src[0] != nullptr) { + std::cerr << "tensor->src[0]=" << tensor->src[0] << " name=" << tensor->src[0]->name << " op=" << ggml_op_name(tensor->src[0]->op) << " type=" << ggml_type_name(tensor->src[0]->type) << " ne0=" << tensor->src[0]->ne[0] << " nb0=" << tensor->src[0]->nb[0] << " ne1=" << tensor->src[0]->ne[1] << " nb1=" << tensor->src[0]->nb[1] << " ne2=" << tensor->src[0]->ne[2] << " nb2=" << tensor->src[0]->nb[2] << " ne3=" << tensor->src[0]->ne[3] << " nb3=" << tensor->src[0]->nb[3] << std::endl; + } + if (tensor->src[1] != nullptr) { + std::cerr << "tensor->src[1]=" << tensor->src[1] << " name=" << tensor->src[1]->name << " op=" << ggml_op_name(tensor->src[1]->op) << " type=" << ggml_type_name(tensor->src[1]->type) << " ne0=" << tensor->src[1]->ne[0] << " nb0=" << tensor->src[1]->nb[0] << " ne1=" << tensor->src[1]->ne[1] << " nb1=" << tensor->src[1]->nb[1] << " ne2=" << tensor->src[1]->ne[2] << " nb2=" << tensor->src[1]->nb[2] << " ne3=" << tensor->src[1]->ne[3] << " nb3=" << tensor->src[1]->nb[3] << std::endl; + } + std::cerr << std::endl << "Result:" << std::endl; + ggml_vk_print_tensor_area(tensor, tensor_data, 5, 5, 0, 0); + std::cerr << std::endl; + std::vector done; + ggml_vk_print_graph_origin(tensor, done); + + if (is_gpu) { + free(tensor_data); + } +} + +void * comp_result; +size_t comp_size; +size_t comp_nb[GGML_MAX_DIMS]; +size_t check_counter = 0; +static void ggml_vk_check_results_0(ggml_tensor * tensor) { + if (tensor->op == GGML_OP_TRANSPOSE) { + return; + } + + check_counter++; + if (!(vk_output_tensor > 0 && vk_output_tensor == check_counter) && check_counter <= vk_skip_checks) { + return; + } + + VK_LOG_DEBUG("ggml_vk_check_results_0(" << tensor->name << ")"); + + ggml_tensor * src0 = tensor->src[0]; + ggml_tensor * src1 = tensor->src[1]; + ggml_tensor * src2 = tensor->src[2]; + ggml_tensor * src3 = tensor->src[3]; + + struct ggml_init_params iparams = { + /*.mem_size =*/ 2ul*1024ul*1024ul*1024ul, + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ false, + }; + + struct ggml_context * ggml_ctx = ggml_init(iparams); + + struct ggml_tensor * src0_clone = nullptr; + struct ggml_tensor * src1_clone = nullptr; + struct ggml_tensor * src2_clone = nullptr; + struct ggml_tensor * src3_clone = nullptr; + struct ggml_tensor * tensor_clone = nullptr; + + size_t src0_size; + size_t src1_size; + size_t src2_size; + size_t src3_size; + + void * src0_buffer = nullptr; + void * src1_buffer = nullptr; + void * src2_buffer = nullptr; + void * src3_buffer = nullptr; + + if (src0 != nullptr) { + src0_clone = ggml_dup_tensor(ggml_ctx, src0); + + src0_size = ggml_nbytes(src0); + + src0_buffer = malloc(src0_size); + src0_clone->data = src0_buffer; + if (ggml_backend_buffer_is_host(src0->buffer)) { + memcpy(src0_clone->data, src0->data, src0_size); + memcpy(src0_clone->nb, src0->nb, sizeof(size_t) * GGML_MAX_DIMS); + } else if (ggml_backend_buffer_is_vk(src0->buffer)) { + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; + vk_buffer& buffer_gpu = buf_ctx->dev_buffer; + uint64_t offset = vk_tensor_offset(src0) + src0->view_offs; + if (!ggml_is_contiguous(src0) && ggml_vk_dim01_contiguous(src0)) { + for (int i3 = 0; i3 < src0->ne[3]; i3++) { + for (int i2 = 0; i2 < src0->ne[2]; i2++) { + const int idx = i3*src0->ne[2] + i2; + ggml_vk_buffer_read(buffer_gpu, offset + idx * src0->nb[2], ((char *)src0_clone->data + idx * src0_clone->nb[2]), src0->ne[1] * src0->nb[1]); + } + } + + src0_clone->nb[0] = src0->nb[0]; + src0_clone->nb[1] = src0->nb[1]; + for (int i = 2; i < GGML_MAX_DIMS; i++) { + src0_clone->nb[i] = src0_clone->nb[i - 1]*src0_clone->ne[i - 1]; + } + } else { + if (offset + src0_size >= buffer_gpu->size) { + src0_size = buffer_gpu->size - offset; + } + ggml_vk_buffer_read(buffer_gpu, offset, src0_clone->data, src0_size); + memcpy(src0_clone->nb, src0->nb, sizeof(size_t) * GGML_MAX_DIMS); + } + } else { + GGML_ABORT("fatal error"); + } + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + ggml_vk_print_tensor(src0, "src0"); + } + } + if (src1 != nullptr) { + src1_clone = ggml_dup_tensor(ggml_ctx, src1); + + src1_size = ggml_nbytes(src1); + + src1_buffer = malloc(src1_size); + src1_clone->data = src1_buffer; + if (ggml_backend_buffer_is_host(src1->buffer)) { + memcpy(src1_clone->data, src1->data, src1_size); + memcpy(src1_clone->nb, src1->nb, sizeof(size_t) * GGML_MAX_DIMS); + } else if (ggml_backend_buffer_is_vk(src1->buffer)) { + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; + vk_buffer& buffer_gpu = buf_ctx->dev_buffer; + uint64_t offset = vk_tensor_offset(src1) + src1->view_offs; + if (!ggml_is_contiguous(src1) && ggml_vk_dim01_contiguous(src1)) { + for (int i3 = 0; i3 < src1->ne[3]; i3++) { + for (int i2 = 0; i2 < src1->ne[2]; i2++) { + const int idx = i3*src1->ne[2] + i2; + ggml_vk_buffer_read(buffer_gpu, offset + idx * src1->nb[2], ((char *)src1_clone->data + idx * src1_clone->nb[2]), src1->ne[1] * src1->nb[1]); + } + } + + src1_clone->nb[0] = src1->nb[0]; + src1_clone->nb[1] = src1->nb[1]; + for (int i = 2; i < GGML_MAX_DIMS; i++) { + src1_clone->nb[i] = src1_clone->nb[i - 1]*src1_clone->ne[i - 1]; + } + } else { + if (offset + src1_size >= buffer_gpu->size) { + src1_size = buffer_gpu->size - offset; + } + ggml_vk_buffer_read(buffer_gpu, offset, src1_clone->data, src1_size); + memcpy(src1_clone->nb, src1->nb, sizeof(size_t) * GGML_MAX_DIMS); + } + } else { + GGML_ABORT("fatal error"); + } + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + ggml_vk_print_tensor(src1, "src1"); + } + } + if (src2 != nullptr) { + src2_clone = ggml_dup_tensor(ggml_ctx, src2); + + src2_size = ggml_nbytes(src2); + + src2_buffer = malloc(src2_size); + src2_clone->data = src2_buffer; + if (ggml_backend_buffer_is_host(src2->buffer)) { + memcpy(src2_clone->data, src2->data, src2_size); + memcpy(src2_clone->nb, src2->nb, sizeof(size_t) * GGML_MAX_DIMS); + } else if (ggml_backend_buffer_is_vk(src2->buffer)) { + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src2->buffer->context; + vk_buffer& buffer_gpu = buf_ctx->dev_buffer; + uint64_t offset = vk_tensor_offset(src2) + src2->view_offs; + if (!ggml_is_contiguous(src2) && ggml_vk_dim01_contiguous(src2)) { + for (int i3 = 0; i3 < src2->ne[3]; i3++) { + for (int i2 = 0; i2 < src2->ne[2]; i2++) { + const int idx = i3*src2->ne[2] + i2; + ggml_vk_buffer_read(buffer_gpu, offset + idx * src2->nb[2], ((char *)src2_clone->data + idx * src2_clone->nb[2]), src2->ne[1] * src2->nb[1]); + } + } + + src2_clone->nb[0] = src2->nb[0]; + src2_clone->nb[1] = src2->nb[1]; + for (int i = 2; i < GGML_MAX_DIMS; i++) { + src2_clone->nb[i] = src2_clone->nb[i - 1]*src2_clone->ne[i - 1]; + } + } else { + if (offset + src2_size >= buffer_gpu->size) { + src2_size = buffer_gpu->size - offset; + } + ggml_vk_buffer_read(buffer_gpu, offset, src2_clone->data, src2_size); + memcpy(src2_clone->nb, src2->nb, sizeof(size_t) * GGML_MAX_DIMS); + } + } else { + GGML_ABORT("fatal error"); + } + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + ggml_vk_print_tensor(src2, "src2"); + } + } + if (src3 != nullptr) { + src3_clone = ggml_dup_tensor(ggml_ctx, src3); + + src3_size = ggml_nbytes(src3); + + src3_buffer = malloc(src3_size); + src3_clone->data = src3_buffer; + if (ggml_backend_buffer_is_host(src3->buffer)) { + memcpy(src3_clone->data, src3->data, src3_size); + memcpy(src3_clone->nb, src3->nb, sizeof(size_t) * GGML_MAX_DIMS); + } else if (ggml_backend_buffer_is_vk(src3->buffer)) { + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src3->buffer->context; + vk_buffer& buffer_gpu = buf_ctx->dev_buffer; + uint64_t offset = vk_tensor_offset(src3) + src3->view_offs; + if (!ggml_is_contiguous(src3) && ggml_vk_dim01_contiguous(src3)) { + for (int i3 = 0; i3 < src3->ne[3]; i3++) { + for (int i2 = 0; i2 < src3->ne[2]; i2++) { + const int idx = i3*src3->ne[2] + i2; + ggml_vk_buffer_read(buffer_gpu, offset + idx * src3->nb[2], ((char *)src3_clone->data + idx * src3_clone->nb[2]), src3->ne[1] * src3->nb[1]); + } + } + + src3_clone->nb[0] = src3->nb[0]; + src3_clone->nb[1] = src3->nb[1]; + for (int i = 2; i < GGML_MAX_DIMS; i++) { + src3_clone->nb[i] = src3_clone->nb[i - 1]*src3_clone->ne[i - 1]; + } + } else { + if (offset + src3_size >= buffer_gpu->size) { + src3_size = buffer_gpu->size - offset; + } + ggml_vk_buffer_read(buffer_gpu, offset, src3_clone->data, src3_size); + memcpy(src3_clone->nb, src3->nb, sizeof(size_t) * GGML_MAX_DIMS); + } + } else { + GGML_ABORT("fatal error"); + } + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + ggml_vk_print_tensor(src3, "src3"); + } + } + + if (tensor->op == GGML_OP_FLASH_ATTN_EXT) { + const float *params = (const float *)tensor->op_params; + tensor_clone = ggml_flash_attn_ext(ggml_ctx, src0_clone, src1_clone, src2_clone, src3_clone, params[0], params[1], params[2]); + } else if (tensor->op == GGML_OP_MUL_MAT) { + tensor_clone = ggml_mul_mat(ggml_ctx, src0_clone, src1_clone); + } else if (tensor->op == GGML_OP_MUL_MAT_ID) { + tensor_clone = ggml_mul_mat_id(ggml_ctx, src0_clone, src1_clone, src2_clone); + } else if (tensor->op == GGML_OP_MUL) { + tensor_clone = ggml_mul(ggml_ctx, src0_clone, src1_clone); + } else if (tensor->op == GGML_OP_DIV) { + tensor_clone = ggml_div(ggml_ctx, src0_clone, src1_clone); + } else if (tensor->op == GGML_OP_CONCAT) { + tensor_clone = ggml_concat(ggml_ctx, src0_clone, src1_clone, *(int *)tensor->op_params); + } else if (tensor->op == GGML_OP_UPSCALE) { + tensor_clone = ggml_upscale_ext(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); + } else if (tensor->op == GGML_OP_SCALE) { + tensor_clone = ggml_scale(ggml_ctx, src0_clone, ((float *)tensor->op_params)[0]); + } else if (tensor->op == GGML_OP_SQR) { + tensor_clone = ggml_sqr(ggml_ctx, src0_clone); + } else if (tensor->op == GGML_OP_SIN) { + tensor_clone = ggml_sin(ggml_ctx, src0_clone); + } else if (tensor->op == GGML_OP_COS) { + tensor_clone = ggml_cos(ggml_ctx, src0_clone); + } else if (tensor->op == GGML_OP_CLAMP) { + tensor_clone = ggml_clamp(ggml_ctx, src0_clone, ((float *)tensor->op_params)[0], ((float *)tensor->op_params)[1]); + } else if (tensor->op == GGML_OP_PAD) { + tensor_clone = ggml_pad(ggml_ctx, src0_clone, tensor->ne[0] - src0_clone->ne[0], tensor->ne[1] - src0_clone->ne[1], tensor->ne[2] - src0_clone->ne[2], tensor->ne[3] - src0_clone->ne[3]); + } else if (tensor->op == GGML_OP_REPEAT) { + tensor_clone = ggml_repeat(ggml_ctx, src0_clone, tensor); + } else if (tensor->op == GGML_OP_ADD) { + tensor_clone = ggml_add(ggml_ctx, src0_clone, src1_clone); + } else if (tensor->op == GGML_OP_ACC) { + tensor_clone = ggml_acc(ggml_ctx, src0_clone, src1_clone, tensor->op_params[0], tensor->op_params[1], tensor->op_params[2], tensor->op_params[3]); + } else if (tensor->op == GGML_OP_NORM) { + tensor_clone = ggml_norm(ggml_ctx, src0_clone, *(float *)tensor->op_params); + } else if (tensor->op == GGML_OP_GROUP_NORM) { + tensor_clone = ggml_group_norm(ggml_ctx, src0_clone, *(int *)tensor->op_params, ((float *)tensor->op_params)[1]); + } else if (tensor->op == GGML_OP_RMS_NORM) { + tensor_clone = ggml_rms_norm(ggml_ctx, src0_clone, *(float *)tensor->op_params); + } else if (tensor->op == GGML_OP_SOFT_MAX) { + if (src1 != nullptr) { + tensor_clone = ggml_soft_max_ext(ggml_ctx, src0_clone, src1_clone, ((float *)tensor->op_params)[0], ((float *)tensor->op_params)[1]); + } else { + tensor_clone = ggml_soft_max(ggml_ctx, src0_clone); + } + } else if (tensor->op == GGML_OP_DIAG_MASK_INF) { + tensor_clone = ggml_diag_mask_inf(ggml_ctx, src0_clone, *(int *)tensor->op_params); + } else if (tensor->op == GGML_OP_ROPE) { + const int n_dims = ((int32_t *) tensor->op_params)[1]; + const int mode = ((int32_t *) tensor->op_params)[2]; + //const int n_ctx_ggml = ((int32_t *) tensor->op_params)[3]; + const int n_ctx_orig_ggml = ((int32_t *) tensor->op_params)[4]; + const float freq_base = ((float *) tensor->op_params)[5]; + const float freq_scale = ((float *) tensor->op_params)[6]; + const float ext_factor = ((float *) tensor->op_params)[7]; + const float attn_factor = ((float *) tensor->op_params)[8]; + const float beta_fast = ((float *) tensor->op_params)[9]; + const float beta_slow = ((float *) tensor->op_params)[10]; + tensor_clone = ggml_rope_ext(ggml_ctx, src0_clone, src1_clone, src2_clone, n_dims, mode, n_ctx_orig_ggml, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + } else if (tensor->op == GGML_OP_UNARY) { + switch (ggml_get_unary_op(tensor)) { + case GGML_UNARY_OP_SILU: + tensor_clone = ggml_silu(ggml_ctx, src0_clone); + break; + case GGML_UNARY_OP_GELU: + tensor_clone = ggml_gelu(ggml_ctx, src0_clone); + break; + case GGML_UNARY_OP_GELU_QUICK: + tensor_clone = ggml_gelu_quick(ggml_ctx, src0_clone); + break; + case GGML_UNARY_OP_RELU: + tensor_clone = ggml_relu(ggml_ctx, src0_clone); + break; + case GGML_UNARY_OP_TANH: + tensor_clone = ggml_tanh(ggml_ctx, src0_clone); + break; + default: + std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl; + GGML_ABORT("fatal error"); + } + } else if (tensor->op == GGML_OP_CPY || tensor->op == GGML_OP_DUP) { + if (src1 == nullptr) { + tensor_clone = ggml_dup(ggml_ctx, src0_clone); + tensor_clone->type = tensor->type; + } else { + tensor_clone = ggml_cpy(ggml_ctx, src0_clone, src1_clone); + } + } else if (tensor->op == GGML_OP_CONT) { + tensor_clone = ggml_cont_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); + } else if (tensor->op == GGML_OP_RESHAPE) { + tensor_clone = ggml_reshape_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); + } else if (tensor->op == GGML_OP_VIEW) { + tensor_clone = ggml_view_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3], tensor->nb[1], tensor->nb[2], tensor->nb[3], ((int32_t *) tensor->op_params)[0]); + } else if (tensor->op == GGML_OP_PERMUTE) { + int32_t * params = (int32_t *)tensor->op_params; + tensor_clone = ggml_permute(ggml_ctx, src0_clone, params[0], params[1], params[2], params[3]); + } else if (tensor->op == GGML_OP_TRANSPOSE) { + tensor_clone = ggml_transpose(ggml_ctx, src0_clone); + } else if (tensor->op == GGML_OP_GET_ROWS) { + tensor_clone = ggml_get_rows(ggml_ctx, src0_clone, src1_clone); + } else if (tensor->op == GGML_OP_ARGSORT) { + tensor_clone = ggml_argsort(ggml_ctx, src0_clone, (ggml_sort_order) *(int *)tensor->op_params); + } else if (tensor->op == GGML_OP_SUM_ROWS) { + tensor_clone = ggml_sum_rows(ggml_ctx, src0_clone); + } else if (tensor->op == GGML_OP_IM2COL) { + const int32_t s0 = tensor->op_params[0]; + const int32_t s1 = tensor->op_params[1]; + const int32_t p0 = tensor->op_params[2]; + const int32_t p1 = tensor->op_params[3]; + const int32_t d0 = tensor->op_params[4]; + const int32_t d1 = tensor->op_params[5]; + + const bool is_2D = tensor->op_params[6] == 1; + tensor_clone = ggml_im2col(ggml_ctx, src0_clone, src1_clone, s0, s1, p0, p1, d0, d1, is_2D, tensor->type); + } else if (tensor->op == GGML_OP_TIMESTEP_EMBEDDING) { + const int32_t dim = tensor->op_params[0]; + const int32_t max_period = tensor->op_params[1]; + tensor_clone = ggml_timestep_embedding(ggml_ctx, src0_clone, dim, max_period); + } else if (tensor->op == GGML_OP_POOL_2D) { + enum ggml_op_pool op = static_cast(tensor->op_params[0]); + const int32_t k0 = tensor->op_params[1]; + const int32_t k1 = tensor->op_params[2]; + const int32_t s0 = tensor->op_params[3]; + const int32_t s1 = tensor->op_params[4]; + const int32_t p0 = tensor->op_params[5]; + const int32_t p1 = tensor->op_params[6]; + + tensor_clone = ggml_pool_2d(ggml_ctx, src0_clone, op, k0, k1, s0, s1, p0, p1); + } else if (tensor->op == GGML_OP_LEAKY_RELU) { + const float * op_params = (const float *)tensor->op_params; + tensor_clone = ggml_leaky_relu(ggml_ctx, src0_clone, op_params[0], false); + } else if (tensor->op == GGML_OP_RWKV_WKV6) { + tensor_clone = ggml_rwkv_wkv6(ggml_ctx, tensor->src[0], tensor->src[1], tensor->src[2], tensor->src[3], + tensor->src[4], tensor->src[5]); + } + else { + std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl; + GGML_ABORT("fatal error"); + } + + ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); + ggml_build_forward_expand(cgraph, tensor_clone); + + ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 8); + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + ggml_vk_print_tensor(tensor_clone, "tensor_clone"); + } + + comp_size = ggml_nbytes(tensor_clone); + + comp_result = malloc(comp_size); + memcpy(comp_result, tensor_clone->data, comp_size); + memcpy(comp_nb, tensor_clone->nb, sizeof(size_t) * GGML_MAX_DIMS); + + if (src0 != nullptr) { + free(src0_buffer); + } + if (src1 != nullptr) { + free(src1_buffer); + } + + ggml_free(ggml_ctx); + + VK_LOG_DEBUG("END ggml_vk_check_results_0(" << tensor->name << ")"); +} + +static void ggml_vk_check_results_1(ggml_tensor * tensor) { + if (tensor->op == GGML_OP_TRANSPOSE) { + return; + } + if (!(vk_output_tensor > 0 && vk_output_tensor == check_counter) && check_counter <= vk_skip_checks) { + return; + } + + VK_LOG_DEBUG("ggml_vk_check_results_1(" << tensor->name << ")"); + + ggml_tensor * src0 = tensor->src[0]; + ggml_tensor * src1 = tensor->src[1]; + ggml_tensor * src2 = tensor->src[2]; + + void * tensor_data = tensor->data; + + if (ggml_backend_buffer_is_vk(tensor->buffer)) { + size_t tensor_size = ggml_nbytes(tensor); + tensor_data = malloc(tensor_size); + + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; + + vk_buffer& buffer_gpu = buf_ctx->dev_buffer; + uint64_t offset = vk_tensor_offset(tensor) + tensor->view_offs; + if (offset + tensor_size >= buffer_gpu->size) { + tensor_size = buffer_gpu->size - offset; + } + + ggml_vk_buffer_read(buffer_gpu, offset, tensor_data, tensor_size); + } + + float first_error_result = -1.0f; + float first_error_correct = -1.0f; + std::array first_error = { -1, -1, -1, -1 }; + double avg_err = 0.0; + size_t counter = 0; + + for (int i3 = 0; i3 < tensor->ne[3]; i3++) { + for (int i2 = 0; i2 < tensor->ne[2]; i2++) { + for (int i1 = 0; i1 < tensor->ne[1]; i1++) { + for (int i0 = 0; i0 < tensor->ne[0]; i0++) { + const bool buffer_size_fit = i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0] < comp_size; + float correct = 0.0f; + float result = 0.0f; + + if (buffer_size_fit) { + if (tensor->type == GGML_TYPE_F32) { + correct = *(float *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0]); + result = *(float *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0]); + } else if (tensor->type == GGML_TYPE_F16) { + correct = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0])); + result = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0])); + } else if (tensor->type == GGML_TYPE_I32) { + correct = *(int32_t *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0]); + result = *(int32_t *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0]); + } else { + std::cerr << "Results check not implemented for type " << ggml_type_name(tensor->type) << std::endl; + } + } else { + std::cerr << "Missing debug code for type " << ggml_type_name(tensor->type) << std::endl; + GGML_ABORT("fatal error"); + } + + if ((std::isnan(correct) != std::isnan(result)) || (std::isinf(correct) != std::isinf(result)) || !buffer_size_fit) { + std::cerr << "ERROR: Invalid value in " << ggml_op_name(tensor->op) << " i3=" << i3 << " i2=" << i2 << " i1=" << i1 << " i0=" << i0 << " result=" << result << " correct=" << correct << " avg_err=" << (avg_err / counter) << std::endl; + std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; + if (src0 != nullptr) { + std::cerr << "src0=" << src0 << " src0->name=" << src0->name << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; + } + if (src1 != nullptr) { + std::cerr << "src1=" << src1 << " src1->name=" << src1->name << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; + } + if (src2 != nullptr) { + std::cerr << "src2=" << src2 << " src2->name=" << src2->name << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; + } + std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; + std::cerr << std::endl << "Result:" << std::endl; + ggml_vk_print_tensor_area(tensor, tensor_data, i0, i1, i2, i3); + std::cerr << std::endl << "Correct:" << std::endl; + ggml_vk_print_tensor_area(tensor, comp_result, i0, i1, i2, i3); + std::cerr << std::endl; + std::vector done; + ggml_vk_print_graph_origin(tensor, done); + GGML_ABORT("fatal error"); + } + if (first_error[0] == -1 && std::fabs(correct - result) > 0.1f) { + first_error[0] = i0; + first_error[1] = i1; + first_error[2] = i2; + first_error[3] = i3; + first_error_result = result; + first_error_correct = correct; + } + + // Special case, value is infinite, avoid NaN result in avg_err + // NaN also appears in results, if both are nan error is 0 + if (!std::isinf(correct) && !std::isinf(result) && !std::isnan(correct) && !std::isnan(result)) { + avg_err += std::fabs(correct - result); + } + counter++; + } + } + } + } + + avg_err /= counter; + + if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { + std::cerr << "TENSOR CHECK: avg_err=" << avg_err << " in " << ggml_op_name(tensor->op) << " (check " << check_counter << ")" << std::endl; + std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; + if (src0 != nullptr) { + std::cerr << "src0=" << src0 << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; + } + if (src1 != nullptr) { + std::cerr << "src1=" << src1 << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; + } + if (src2 != nullptr) { + std::cerr << "src2=" << src2 << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; + } + std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; + std::cerr << std::endl << "Result:" << std::endl; + ggml_vk_print_tensor_area(tensor, tensor_data, 5, 5, 0, 0); + std::cerr << std::endl << "Correct:" << std::endl; + ggml_vk_print_tensor_area(tensor, comp_result, 5, 5, 0, 0); + std::cerr << std::endl; + std::vector done; + ggml_vk_print_graph_origin(tensor, done); + } + + if (avg_err > 0.05 || std::isnan(avg_err)) { + std::cerr << "ERROR: avg_err=" << avg_err << " in " << ggml_op_name(tensor->op) << " (check " << check_counter << ")" << std::endl; + std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; + if (src0 != nullptr) { + std::cerr << "src0=" << src0 << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; + } + if (src1 != nullptr) { + std::cerr << "src1=" << src1 << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; + } + if (src2 != nullptr) { + std::cerr << "src2=" << src2 << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; + } + std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; + std::cerr << std::endl << "Result:" << std::endl; + ggml_vk_print_tensor_area(tensor, tensor_data, first_error[0], first_error[1], first_error[2], first_error[3]); + std::cerr << std::endl << "Correct:" << std::endl; + ggml_vk_print_tensor_area(tensor, comp_result, first_error[0], first_error[1], first_error[2], first_error[3]); + std::cerr << std::endl; + std::vector done; + ggml_vk_print_graph_origin(tensor, done); + GGML_ABORT("fatal error"); + } else { + std::cerr << check_counter << " " << tensor->name << " op=" << ggml_op_name(tensor->op) << " avg_err=" << avg_err << std::endl; + } + + free(comp_result); + comp_result = nullptr; + comp_size = 0; + + if (ggml_backend_buffer_is_vk(tensor->buffer)) { + free(tensor_data); + } + + VK_LOG_DEBUG("END ggml_vk_check_results_1(" << tensor->name << ")"); +} +#endif + +GGML_BACKEND_DL_IMPL(ggml_backend_vk_reg) diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt new file mode 100644 index 000000000..bd0c74cb1 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt @@ -0,0 +1,9 @@ +find_package (Threads REQUIRED) +find_package(Vulkan COMPONENTS glslc REQUIRED) + +set(TARGET vulkan-shaders-gen) +add_executable(${TARGET} vulkan-shaders-gen.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_compile_features(${TARGET} PRIVATE cxx_std_17) +target_link_libraries(vulkan-shaders-gen PUBLIC Threads::Threads) +target_link_libraries(vulkan-shaders-gen PRIVATE Vulkan::Vulkan) diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp new file mode 100644 index 000000000..d896f1ef0 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp @@ -0,0 +1,29 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = gl_GlobalInvocationID.x; + if (idx >= p.ne) { + return; + } + + const uint offset = p.param3; + const uint src1_i = idx - offset; + const uint oz = src1_i / p.nb02; + const uint oy = (src1_i - (oz * p.nb02)) / p.nb01; + const uint ox = src1_i % p.nb01; + + uint i00, i01, i02, i03; + get_indices(idx, i00, i01, i02, i03); + + if (ox < p.ne10 && oy < p.ne11 && oz < p.ne12) { + data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) + FLOAT_TYPE(data_b[get_boffset() + ox + oy * p.ne10 + oz * p.ne10 * p.ne11])); + } else { + data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)])); + } +} + diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp new file mode 100644 index 000000000..2b4085c4f --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp @@ -0,0 +1,29 @@ +#version 450 + +#extension GL_EXT_shader_16bit_storage : require + +#include "types.comp" +#include "generic_binary_head.comp" + +const uint num_threads = 256; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + uint idx = get_idx(); + + // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation + const uint num_iter = 2; + + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (idx >= p.ne) { + continue; + } + uint i00, i01, i02, i03; + get_indices(idx, i00, i01, i02, i03); + + data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) + FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); + + idx += num_threads; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp new file mode 100644 index 000000000..d4fa45b1e --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp @@ -0,0 +1,69 @@ +#version 450 + +#include "types.comp" + +#define BLOCK_SIZE 1024 +#define ASC 0 + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) buffer D {int data_d[];}; + +layout (push_constant) uniform parameter { + uint ncols; + uint ncols_pad; + uint order; +} p; + +shared int dst_row[BLOCK_SIZE]; + +void swap(uint idx0, uint idx1) { + int tmp = dst_row[idx0]; + dst_row[idx0] = dst_row[idx1]; + dst_row[idx1] = tmp; +} + +void main() { + // bitonic sort + const int col = int(gl_LocalInvocationID.x); + const uint row = gl_WorkGroupID.y; + + const uint row_offset = row * p.ncols; + + // initialize indices + if (col < p.ncols_pad) { + dst_row[col] = col; + } + barrier(); + + for (uint k = 2; k <= p.ncols_pad; k *= 2) { + for (uint j = k / 2; j > 0; j /= 2) { + const uint ixj = col ^ j; + if (col < p.ncols_pad && ixj > col) { + if ((col & k) == 0) { + if (dst_row[col] >= p.ncols || + (dst_row[ixj] < p.ncols && (p.order == ASC ? + data_a[row_offset + dst_row[col]] > data_a[row_offset + dst_row[ixj]] : + data_a[row_offset + dst_row[col]] < data_a[row_offset + dst_row[ixj]])) + ) { + swap(col, ixj); + } + } else { + if (dst_row[ixj] >= p.ncols || + (dst_row[col] < p.ncols && (p.order == ASC ? + data_a[row_offset + dst_row[col]] < data_a[row_offset + dst_row[ixj]] : + data_a[row_offset + dst_row[col]] > data_a[row_offset + dst_row[ixj]])) + ) { + swap(col, ixj); + } + } + } + barrier(); + } + } + + if (col < p.ncols) { + data_d[row_offset + col] = dst_row[col]; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp new file mode 100644 index 000000000..1e5cb8dae --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp @@ -0,0 +1,17 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + + const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(val < p.param1 ? p.param1 : (val > p.param2 ? p.param2 : val)); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp new file mode 100644 index 000000000..9ee2f1fae --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp @@ -0,0 +1,41 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + const int dim = p.param3; + + if (idx >= p.ne) { + return; + } + + const uint i3 = idx / (p.ne22*p.ne21*p.ne20); + const uint i3_offset = i3 * p.ne22*p.ne21*p.ne20; + const uint i2 = (idx - i3_offset) / (p.ne21*p.ne20); + const uint i2_offset = i2*p.ne21*p.ne20; + const uint i1 = (idx - i3_offset - i2_offset) / p.ne20; + const uint i0 = idx - i3_offset - i2_offset - i1*p.ne20; + + uint o[4] = {0, 0, 0, 0}; + o[dim] = dim == 0 ? p.ne00 : (dim == 1 ? p.ne01 : (dim == 2 ? p.ne02 : p.ne03)); + + const uint src0_idx = i3*p.nb03 + i2*p.nb02 + i1*p.nb01 + i0*p.nb00; + const uint src1_idx = (i3 - o[3])*p.nb13 + (i2 - o[2])*p.nb12 + (i1 - o[1])*p.nb11 + (i0 - o[0])*p.nb10; + const uint dst_idx = i3*p.nb23 + i2*p.nb22 + i1*p.nb21 + i0*p.nb20; + + const bool is_src0 = i0 < p.ne00 && i1 < p.ne01 && i2 < p.ne02 && i3 < p.ne03; + +#ifndef OPTIMIZATION_ERROR_WORKAROUND + data_d[get_doffset() + dst_idx] = D_TYPE(is_src0 ? data_a[get_aoffset() + src0_idx] : data_b[get_boffset() + src1_idx]); +#else + if (is_src0) { + data_d[get_doffset() + dst_idx] = data_a[get_aoffset() + src0_idx]; + } else { + data_d[get_doffset() + dst_idx] = data_b[get_boffset() + src1_idx]; + } +#endif +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp new file mode 100644 index 000000000..dd828c232 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp @@ -0,0 +1,42 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +#extension GL_EXT_control_flow_attributes : require + +const uint num_threads = 128; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + uint idx = get_idx(); + + // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation + const uint num_iter = 4; + + // fast path for when all four iterations are in-bounds + if (idx + (num_iter-1)*num_threads < p.ne) { + [[unroll]] for (uint i = 0; i < num_iter; ++i) { +#ifndef OPTIMIZATION_ERROR_WORKAROUND + data_d[get_doffset() + idx] = D_TYPE(data_a[get_aoffset() + idx]); +#else + data_d[get_doffset() + idx] = data_a[get_aoffset() + idx]; +#endif + idx += num_threads; + } + } else { + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (idx >= p.ne) { + continue; + } + +#ifndef OPTIMIZATION_ERROR_WORKAROUND + data_d[get_doffset() + idx] = D_TYPE(data_a[get_aoffset() + idx]); +#else + data_d[get_doffset() + idx] = data_a[get_aoffset() + idx]; +#endif + idx += num_threads; + } + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp new file mode 100644 index 000000000..29c906494 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp @@ -0,0 +1,20 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + +#ifndef OPTIMIZATION_ERROR_WORKAROUND + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(data_a[get_aoffset() + src0_idx(idx)]); +#else + data_d[get_doffset() + dst_idx(idx)] = data_a[get_aoffset() + src0_idx(idx)]; +#endif +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp new file mode 100644 index 000000000..0b8d02f58 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp @@ -0,0 +1,17 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + + const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(cos(val)); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp new file mode 100644 index 000000000..a4d3fca55 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp @@ -0,0 +1,20 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {float data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_GlobalInvocationID.x * 16; + + if (i >= p.nel) { + return; + } + + [[unroll]] for (uint l = 0; l < 16; l++) { + data_b[i + l] = D_TYPE(data_a[i + l]); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp new file mode 100644 index 000000000..91bb8f8db --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp @@ -0,0 +1,118 @@ +#if !defined(DATA_A_F32) && !defined(DATA_A_F16) +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#endif + +#include "types.comp" + +#if defined(A_TYPE_PACKED16) +layout (binding = 0) readonly buffer A_PACKED16 {A_TYPE_PACKED16 data_a_packed16[];}; +#endif +#if defined(A_TYPE_PACKED32) +layout (binding = 0) readonly buffer A_PACKED32 {A_TYPE_PACKED32 data_a_packed32[];}; +#endif + +#if defined(DATA_A_F32) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]); +} +#endif + +#if defined(DATA_A_F16) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]); +} +#endif + +#if defined(DATA_A_Q4_0) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + return (vec2(vui & 0xF, vui >> 4) - 8.0f); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); + return (vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, vui >> 12) - 8.0f); +} +#endif + +#if defined(DATA_A_Q4_1) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + return vec2(vui & 0xF, vui >> 4); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); + return vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, vui >> 12); +} +#endif + +#if defined(DATA_A_Q5_0) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint uint_qh = uint(data_a[a_offset + ib].qh[1]) << 16 | data_a[a_offset + ib].qh[0]; + const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + return (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint uint_qh = uint(data_a_packed16[a_offset + ib].qh[1]) << 16 | data_a_packed16[a_offset + ib].qh[0]; + const ivec2 qh0 = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const ivec2 qh1 = ivec2(((uint_qh >> (iqs + 1)) << 4) & 0x10, (uint_qh >> (iqs + 13)) & 0x10); + const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); + return (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f); +} +#endif + +#if defined(DATA_A_Q5_1) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint uint_qh = data_a[a_offset + ib].qh; + const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + return vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint uint_qh = data_a_packed16[a_offset + ib].qh; + const ivec2 qh0 = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const ivec2 qh1 = ivec2(((uint_qh >> (iqs + 1)) << 4) & 0x10, (uint_qh >> (iqs + 13)) & 0x10); + const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); + return vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y); +} +#endif + +#if defined(DATA_A_Q8_0) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + return vec2(int(data_a[a_offset + ib].qs[iqs]), int(data_a[a_offset + ib].qs[iqs + 1])); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + uint32_t v0 = data_a_packed16[a_offset + ib].qs[iqs/2]; + uint32_t v1 = data_a_packed16[a_offset + ib].qs[iqs/2 + 1]; + return vec4(int8_t(v0 & 0xFF), int8_t(v0 >> 8), int8_t(v1 & 0xFF), int8_t(v1 >> 8)); +} +#endif + +#if defined(DATA_A_IQ4_NL) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + return vec2(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[vui >> 4]); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); + return vec4(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[(vui >> 4) & 0xF], kvalues_iq4nl[(vui >> 8) & 0xF], kvalues_iq4nl[vui >> 12]); +} +#endif + +#if defined(DATA_A_F32) || defined(DATA_A_F16) +vec2 get_dm(uint ib, uint a_offset) { + return vec2(0, 0); +} +#endif + +#if defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ4_NL) +vec2 get_dm(uint ib, uint a_offset) { + return vec2(float(data_a[a_offset + ib].d), 0); +} +#endif + +#if defined(DATA_A_Q4_1) || defined(DATA_A_Q5_1) +vec2 get_dm(uint ib, uint a_offset) { + return vec2(float(data_a[a_offset + ib].d), float(data_a[a_offset + ib].m)); +} +#endif diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp new file mode 100644 index 000000000..94b78598e --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp @@ -0,0 +1,325 @@ + +#include "types.comp" + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ4_0 { + block_q4_0_packed16 block; +}; + +float16_t dequantFuncQ4_0(const in decodeBufQ4_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const uint idx = coordInBlock[1]; + const uint shift = (idx & 0x10) >> 2; + uint32_t qs = uint32_t(bl.block.qs[(idx & 0xE) >> 1]); + qs >>= shift; + qs &= 0x0F0F; + qs = unpack8(qs)[idx & 1]; + float16_t ret = (float16_t(qs) - float16_t(8)) * d; + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 4) buffer decodeBufQ4_1 { + block_q4_1 block; +}; + +float16_t dequantFuncQ4_1(const in decodeBufQ4_1 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const float16_t m = bl.block.m; + const uint idx = coordInBlock[1]; + const uint iqs = idx & 0xF; + const uint shift = (idx & 0x10) >> 2; + uint32_t qs = bl.block.qs[iqs]; + qs >>= shift; + qs &= 0xF; + float16_t ret = float16_t(qs) * d + m; + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ5_0 { + block_q5_0 block; +}; + +float16_t dequantFuncQ5_0(const in decodeBufQ5_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const uint idx = coordInBlock[1]; + const uint iqs = idx & 0xF; + + const uint uint_qh = uint(bl.block.qh[1]) << 16 | bl.block.qh[0]; + const uint qh = ((uint_qh >> idx) << 4) & 0x10; + + const uint shift = (idx & 0x10) >> 2; + uint32_t qs = bl.block.qs[iqs]; + qs >>= shift; + qs &= 0xF; + + float16_t ret = (float16_t(qs | qh) - float16_t(16)) * d; + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 8) buffer decodeBufQ5_1 { + block_q5_1 block; +}; + +float16_t dequantFuncQ5_1(const in decodeBufQ5_1 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const float16_t m = bl.block.m; + const uint idx = coordInBlock[1]; + const uint iqs = idx & 0xF; + + const uint uint_qh = bl.block.qh; + const uint qh = ((uint_qh >> idx) << 4) & 0x10; + + const uint shift = (idx & 0x10) >> 2; + uint32_t qs = bl.block.qs[iqs]; + qs >>= shift; + qs &= 0xF; + + float16_t ret = float16_t(qs | qh) * d + m; + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ8_0 { + block_q8_0_packed16 block; +}; + +float16_t dequantFuncQ8_0(const in decodeBufQ8_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const uint idx = coordInBlock[1]; + const uint iqs = idx; + + // Load 16b and select the byte for this element + int32_t qs = unpack8(int32_t(bl.block.qs[(iqs & 0x1E) >> 1]))[iqs & 1]; + float16_t ret = float16_t(qs) * d; + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 4) buffer decodeBufQ2_K { + block_q2_K block; +}; + +float16_t dequantFuncQ2_K(const in decodeBufQ2_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const f16vec2 d = bl.block.d; + const uint idx = coordInBlock[1]; + const uint iqs = idx; + + const uint qsi = (iqs / 128) * 32 + (iqs % 32); // 0..31 + const uint scalesi = iqs / 16; // 0..15 + const uint qsshift = ((iqs % 128) / 32) * 2; // 0,2,4,6 + + uint32_t qs = bl.block.qs[qsi]; + const uint scales = bl.block.scales[scalesi]; + float16_t ret = d.x * float16_t(scales & 0xF) * float16_t((qs >> qsshift) & 3) - d.y * float16_t(scales >> 4); + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ3_K { + block_q3_K block; +}; + +float16_t dequantFuncQ3_K(const in decodeBufQ3_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint idx = coordInBlock[1]; + const uint iqs = idx; + + const uint n = iqs / 128; // 0,1 + const uint qsi = n * 32 + (iqs % 32); // 0..63 + const uint hmi = (iqs % 32); // 0..31 + const uint j = (iqs % 128) / 8; // 0..15 + const uint is = iqs / 16; // 0..15 + const uint halfsplit = ((iqs % 128) / 32); // 0,1,2,3 + const uint qsshift = halfsplit * 2; // 0,2,4,6 + const uint m = 1 << (4 * n + halfsplit); // 1,2,4,8,16,32,64,128 + + uint32_t scaleidx0 = (is < 8) ? is : (is-8); + uint32_t scaleidx0shift = (is < 8) ? 0 : 4; + uint32_t scaleidx1 = is + 8 - (is/4)*4; + uint32_t scaleidx1shift = (is/4)*2; + + const int8_t us = int8_t(((bl.block.scales[scaleidx0] >> scaleidx0shift) & 0xF) | (((bl.block.scales[scaleidx1] >> scaleidx1shift) & 3) << 4)); + + const float16_t dl = bl.block.d * float16_t(us - 32); + + float16_t ret = dl * float16_t(int8_t((bl.block.qs[qsi ] >> qsshift) & 3) - (((bl.block.hmask[hmi ] & m) != 0) ? 0 : 4)); + + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ4_K { + block_q4_K block; +}; + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ4_K_packed16 { + block_q4_K_packed16 block; +}; + +float16_t dequantFuncQ4_K(const in decodeBufQ4_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + decodeBufQ4_K_packed16 bl16 = decodeBufQ4_K_packed16(bl); + const uint idx = coordInBlock[1]; + + const uint b = (idx & 0x20) >> 5; // 0,1 + const uint is = (idx & 0xE0) >> 5; // 0..7 + + const f16vec2 loadd = bl.block.d; + + uint32_t sc; + uint32_t mbyte; + + uint32_t scidx0 = (is < 4) ? is : (is + 4); + uint32_t scidx1 = (is < 4) ? is : (is - 4); + uint32_t scidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint32_t scidxshift1 = (is < 4) ? 0 : 2; + uint32_t mbidx0 = is + 4; + uint32_t mbidx1 = (is < 4) ? is + 4 : is; + uint32_t mbidxmask0 = (is < 4) ? 0xF : 0xF0; + uint32_t mbidxshift0 = (is < 4) ? 0 : 4; + uint32_t mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint32_t mbidxshift1 = (is < 4) ? 0 : 2; + + sc = uint8_t((bl.block.scales[scidx0] & 0xF) | ((bl.block.scales[scidx1] & scidxmask1) >> scidxshift1)); + mbyte = uint8_t(((bl.block.scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((bl.block.scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const float16_t d = loadd.x * float16_t(sc); + const float16_t m = loadd.y * float16_t(mbyte); + + uint qs = uint32_t(bl16.block.qs[((idx & 0xC0) >> 2) + ((idx & 0x1E) >> 1)]); + qs = (qs >> (b * 4)) & 0x0F0F; + qs = unpack8(qs)[idx & 1]; + + float16_t ret = d * float16_t(qs) - m; + + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ5_K { + block_q5_K block; +}; + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ5_K_packed16 { + block_q5_K_packed16 block; +}; + +float16_t dequantFuncQ5_K(const in decodeBufQ5_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + decodeBufQ5_K_packed16 bl16 = decodeBufQ5_K_packed16(bl); + const uint idx = coordInBlock[1]; + + const uint b = (idx & 0x20) >> 5; // 0,1 + const uint is = (idx & 0xE0) >> 5; // 0..7 + + const uint32_t hm = 0x0101 << is; + + const f16vec2 loadd = bl.block.d; + + uint32_t sc; + uint32_t mbyte; + + uint32_t scidx0 = (is < 4) ? is : (is + 4); + uint32_t scidx1 = (is < 4) ? is : (is - 4); + uint32_t scidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint32_t scidxshift1 = (is < 4) ? 0 : 2; + uint32_t mbidx0 = is + 4; + uint32_t mbidx1 = (is < 4) ? is + 4 : is; + uint32_t mbidxmask0 = (is < 4) ? 0xF : 0xF0; + uint32_t mbidxshift0 = (is < 4) ? 0 : 4; + uint32_t mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint32_t mbidxshift1 = (is < 4) ? 0 : 2; + + sc = uint8_t((bl.block.scales[scidx0] & 0xF) | ((bl.block.scales[scidx1] & scidxmask1) >> scidxshift1)); + mbyte = uint8_t(((bl.block.scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((bl.block.scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const float16_t d = loadd.x * float16_t(sc); + const float16_t m = loadd.y * float16_t(mbyte); + + uint qh = uint32_t(bl16.block.qh[(idx & 0x1E) >> 1]); + qh = qh & hm; + qh = unpack8(qh)[idx & 1]; + + uint qs = uint32_t(bl16.block.qs[((idx & 0xC0) >> 2) + ((idx & 0x1E) >> 1)]); + qs = (qs >> (b * 4)) & 0x0F0F; + qs = unpack8(qs)[idx & 1]; + + float16_t ret = d * (float16_t(qs) + (qh != 0 ? float16_t(16) : float16_t(0))) - m; + + return ret; +} + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ6_K { + block_q6_K block; +}; + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ6_K_packed16 { + block_q6_K_packed16 block; +}; + +float16_t dequantFuncQ6_K(const in decodeBufQ6_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + decodeBufQ6_K_packed16 bl16 = decodeBufQ6_K_packed16(bl); + const uint idx = coordInBlock[1]; + + const uint b = (idx & 0x40) >> 6; // 0,1 + const uint qhshift = (idx & 0x60) >> 4; // 0,2,4,6 + const uint is = (idx & 0xF0) >> 4; // 0..15 + + const float16_t dscale = bl.block.d * float16_t(bl.block.scales[is]); + + uint ql = uint32_t(bl16.block.ql[((idx & 0x80) >> 2) + ((idx & 0x3E) >> 1)]); + ql = (ql >> (b * 4)) & 0x0F0F; + + uint qh = uint32_t(bl16.block.qh[((idx & 0x80) >> 3) + ((idx & 0x1E) >> 1)]); + qh = ((qh >> qhshift) & 0x0303) << 4; + + int q = unpack8(ql | qh)[idx & 1]; + + float16_t ret = dscale * float16_t(q - 32); + + return ret; +} + +#if defined(DATA_A_IQ4_NL) +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufIQ4_NL { + block_iq4_nl block; +}; + +float16_t dequantFuncIQ4_NL(const in decodeBufIQ4_NL bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const float16_t d = bl.block.d; + const uint idx = coordInBlock[1]; + const uint iqs = idx & 0xF; + const uint shift = (idx & 0x10) >> 2; + uint32_t qs = bl.block.qs[iqs]; + qs >>= shift; + qs &= 0xF; + float16_t ret = float16_t(kvalues_iq4nl[qs]) * d; + return ret; +} +#endif + +#if defined(DATA_A_Q4_0) +#define dequantFuncA dequantFuncQ4_0 +#elif defined(DATA_A_Q4_1) +#define dequantFuncA dequantFuncQ4_1 +#elif defined(DATA_A_Q5_0) +#define dequantFuncA dequantFuncQ5_0 +#elif defined(DATA_A_Q5_1) +#define dequantFuncA dequantFuncQ5_1 +#elif defined(DATA_A_Q8_0) +#define dequantFuncA dequantFuncQ8_0 +#elif defined(DATA_A_Q2_K) +#define dequantFuncA dequantFuncQ2_K +#elif defined(DATA_A_Q3_K) +#define dequantFuncA dequantFuncQ3_K +#elif defined(DATA_A_Q4_K) +#define dequantFuncA dequantFuncQ4_K +#elif defined(DATA_A_Q5_K) +#define dequantFuncA dequantFuncQ5_K +#elif defined(DATA_A_Q6_K) +#define dequantFuncA dequantFuncQ6_K +#elif defined(DATA_A_IQ4_NL) +#define dequantFuncA dequantFuncIQ4_NL +#endif diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp new file mode 100644 index 000000000..8d806435b --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp @@ -0,0 +1,13 @@ +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require + +layout (push_constant) uniform parameter +{ + uint M; + uint K; + uint stride_a; + uint stride_b; + uint nel; +} p; + +#include "types.comp" diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp new file mode 100644 index 000000000..8de14fc03 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp @@ -0,0 +1,32 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_iq4_nl data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + init_iq4nl_shmem(); + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint q_idx = 8*il; + const uint b_idx = 1024*i + 32*ir + q_idx; + + const float d = float(data_a[ib].d); + + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_b[b_idx + l + 0] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] & 0xF]); + data_b[b_idx + l + 16] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] >> 4]); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp new file mode 100644 index 000000000..157154af3 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp @@ -0,0 +1,34 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { + const uint i = gl_WorkGroupID.x * 256 + wgy; + if (i >= p.M * p.K / QUANT_K) { + return; + } + + const uint tid = gl_LocalInvocationID.x; + const uint ip = tid / 32; + const uint il = tid - 32 * ip; + const uint is = 8 * ip + il / 16; + + const uint y_idx = i * QUANT_K + 128 * ip + il; + + const uint ql_idx = 32 * ip + il; + const uint8_t qs = data_a[i].qs[32 * ip + il]; + + FLOAT_TYPE dall = FLOAT_TYPE(data_a[i].d.x); + FLOAT_TYPE dmin = FLOAT_TYPE(data_a[i].d.y); + data_b[y_idx + 0] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+0] & 0xF) * ((qs >> 0) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+0] >> 4)); + data_b[y_idx + 32] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+2] & 0xF) * ((qs >> 2) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+2] >> 4)); + data_b[y_idx + 64] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+4] & 0xF) * ((qs >> 4) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+4] >> 4)); + data_b[y_idx + 96] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+6] & 0xF) * ((qs >> 6) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+6] >> 4)); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp new file mode 100644 index 000000000..c17dd0d99 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp @@ -0,0 +1,42 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { + const uint i = uint(gl_WorkGroupID.x * 256 + wgy); + if (i >= p.M * p.K / QUANT_K) { + return; + } + + const uint r = gl_LocalInvocationID.x / 4; + const uint tid = r / 2; + const uint is0 = r % 2; + const uint l0 = 16 * is0 + 4 * (gl_LocalInvocationID.x % 4); + const uint n = tid / 4; + const uint j = tid - 4*n; + + const uint8_t m = uint8_t(1 << (4*n + j)); + const uint is = 8*n + 2*j + is0; + const uint shift = 2*j; + + const int8_t us = int8_t(is < 4 ? (data_a[i].scales[is-0] & 0xF) | (((data_a[i].scales[is+8] >> 0) & 3) << 4) : + is < 8 ? (data_a[i].scales[is-0] & 0xF) | (((data_a[i].scales[is+4] >> 2) & 3) << 4) : + is < 12 ? (data_a[i].scales[is-8] >> 4) | (((data_a[i].scales[is+0] >> 4) & 3) << 4) : + (data_a[i].scales[is-8] >> 4) | (((data_a[i].scales[is-4] >> 6) & 3) << 4)); + const FLOAT_TYPE d_all = FLOAT_TYPE(data_a[i].d); + const FLOAT_TYPE dl = d_all * FLOAT_TYPE(us - 32); + + const uint y_idx = i * QUANT_K + 128 * n + 32 * j; + const uint qs_idx = 32*n; + + for (uint l = l0; l < l0 + 4; ++l) { + data_b[y_idx + l] = D_TYPE(dl * FLOAT_TYPE(int8_t((data_a[i].qs[qs_idx + l] >> shift) & 3) - (((data_a[i].hmask[l] & m) != 0) ? 0 : 4))); + } + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp new file mode 100644 index 000000000..408185327 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp @@ -0,0 +1,30 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_q4_0 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint q_idx = 8*il; + const uint b_idx = 1024*i + 32*ir + q_idx; + + const float d = float(data_a[ib].d); + + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_b[b_idx + l + 0] = D_TYPE(d * ((data_a[ib].qs[q_idx + l] & 0xF) - 8.0f)); + data_b[b_idx + l + 16] = D_TYPE(d * ((data_a[ib].qs[q_idx + l] >> 4) - 8.0f)); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp new file mode 100644 index 000000000..2f27eee68 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp @@ -0,0 +1,32 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_q4_1 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint b_idx = 1024*i + 32*ir + 8*il; + + const float d = float(data_a[ib].d); + const float m = float(data_a[ib].m); + + const uint q_idx = 8*il; + + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_b[b_idx + l + 0] = D_TYPE(d * (data_a[ib].qs[q_idx + l] & 0xF) + m); + data_b[b_idx + l + 16] = D_TYPE(d * (data_a[ib].qs[q_idx + l] >> 4) + m); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp new file mode 100644 index 000000000..987f113a3 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp @@ -0,0 +1,68 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { + const uint ib = gl_WorkGroupID.x * 256 + wgy; + if (ib >= p.M * p.K / QUANT_K) { + return; + } + + const uint tid = gl_LocalInvocationID.x; + const uint il = tid / 8; + const uint ir = tid % 8; + const uint is = 2 * il; + const uint n = 4; + + const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib].d.x); + const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib].d.y); + + const uint y_idx = ib * QUANT_K + 64 * il + n * ir; + const uint qs_idx = 32*il + n * ir; + + uint scidx0 = (is < 4) ? is : (is + 4); + uint scidx1 = (is < 4) ? is : (is - 4); + uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint scidxshift1 = (is < 4) ? 0 : 2; + uint mbidx0 = is + 4; + uint mbidx1 = (is < 4) ? is + 4 : is; + uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; + uint mbidxshift0 = (is < 4) ? 0 : 4; + uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint mbidxshift1 = (is < 4) ? 0 : 2; + + uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const FLOAT_TYPE d1 = dall * sc; + const FLOAT_TYPE m1 = dmin * mbyte; + + scidx0 = (is < 4) ? is + 1 : (is + 5); + scidx1 = (is < 4) ? is + 1 : (is - 3); + scidxmask1 = (is < 4) ? 0x30 : 0xC0; + scidxshift1 = (is < 4) ? 0 : 2; + mbidx0 = is + 5; + mbidx1 = (is < 4) ? is + 5 : is + 1; + mbidxmask0 = (is < 4) ? 0xF : 0xF0; + mbidxshift0 = (is < 4) ? 0 : 4; + mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + mbidxshift1 = (is < 4) ? 0 : 2; + + sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const FLOAT_TYPE d2 = dall * sc; + const FLOAT_TYPE m2 = dmin * mbyte; + + [[unroll]] for (uint l = 0; l < n; ++l) { + data_b[y_idx + l ] = D_TYPE(d1 * FLOAT_TYPE(data_a[ib].qs[qs_idx + l] & 0xF) - m1); + data_b[y_idx + l + 32] = D_TYPE(d2 * FLOAT_TYPE(data_a[ib].qs[qs_idx + l] >> 4) - m2); + } + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp new file mode 100644 index 000000000..b20b80529 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp @@ -0,0 +1,34 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_q5_0 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint b_idx = 1024*i + 32*ir + 8*il; + + const float d = float(data_a[ib].d); + const uint qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0]; + + const uint q_idx = 8*il; + + [[unroll]] for (uint l = 0; l < 8; ++l) { + const uint iqs = q_idx + l; + const uint vui = uint(data_a[ib].qs[iqs]); + data_b[b_idx + l + 0] = D_TYPE(d * (((vui & 0xF) | (((qh >> iqs) << 4) & 0x10)) - 16.0f)); + data_b[b_idx + l + 16] = D_TYPE(d * (((vui >> 4) | ((qh >> (iqs + 12)) & 0x10)) - 16.0f)); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp new file mode 100644 index 000000000..dc59fe3b7 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp @@ -0,0 +1,35 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_q5_1 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint b_idx = 1024*i + 32*ir + 8*il; + + const float d = float(data_a[ib].d); + const float m = float(data_a[ib].m); + const uint qh = data_a[ib].qh; + + const uint q_idx = 8*il; + + [[unroll]] for (uint l = 0; l < 8; ++l) { + const uint iqs = q_idx + l; + const uint vui = uint(data_a[ib].qs[iqs]); + data_b[b_idx + l + 0] = D_TYPE(d * (((vui & 0xF) | (((qh >> iqs) << 4) & 0x10))) + m); + data_b[b_idx + l + 16] = D_TYPE(d * (((vui >> 4) | ((qh >> (iqs + 12)) & 0x10))) + m); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp new file mode 100644 index 000000000..6db5403b6 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp @@ -0,0 +1,70 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { + const uint ib = gl_WorkGroupID.x * 256 + wgy; + if (ib >= p.M * p.K / QUANT_K) { + return; + } + + const uint tid = gl_LocalInvocationID.x; + const uint il = tid / 16; + const uint ir = tid % 16; + const uint is = 2 * il; + + const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib].d.x); + const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib].d.y); + + const uint y_idx = ib * QUANT_K + 64 * il + 2 * ir; + const uint qs_idx = 32*il + 2 * ir; + const uint qh_idx = 2 * ir; + + uint scidx0 = (is < 4) ? is : (is + 4); + uint scidx1 = (is < 4) ? is : (is - 4); + uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint scidxshift1 = (is < 4) ? 0 : 2; + uint mbidx0 = is + 4; + uint mbidx1 = (is < 4) ? is + 4 : is; + uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; + uint mbidxshift0 = (is < 4) ? 0 : 4; + uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + uint mbidxshift1 = (is < 4) ? 0 : 2; + + uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const FLOAT_TYPE d1 = dall * sc; + const FLOAT_TYPE m1 = dmin * mbyte; + + scidx0 = (is < 4) ? is + 1 : (is + 5); + scidx1 = (is < 4) ? is + 1 : (is - 3); + scidxmask1 = (is < 4) ? 0x30 : 0xC0; + scidxshift1 = (is < 4) ? 0 : 2; + mbidx0 = is + 5; + mbidx1 = (is < 4) ? is + 5 : is + 1; + mbidxmask0 = (is < 4) ? 0xF : 0xF0; + mbidxshift0 = (is < 4) ? 0 : 4; + mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + mbidxshift1 = (is < 4) ? 0 : 2; + + sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const FLOAT_TYPE d2 = dall * sc; + const FLOAT_TYPE m2 = dmin * mbyte; + + const uint8_t hm1 = uint8_t(1 << (2 * il )); + const uint8_t hm2 = uint8_t(1 << (2 * il + 1)); + data_b[y_idx ] = D_TYPE(d1 * FLOAT_TYPE((data_a[ib].qs[qs_idx ] & 0xF) + (((data_a[ib].qh[qh_idx ] & hm1) != 0) ? 16 : 0)) - m1); + data_b[y_idx + 1] = D_TYPE(d1 * FLOAT_TYPE((data_a[ib].qs[qs_idx + 1] & 0xF) + (((data_a[ib].qh[qh_idx + 1] & hm1) != 0) ? 16 : 0)) - m1); + data_b[y_idx + 32] = D_TYPE(d2 * FLOAT_TYPE((data_a[ib].qs[qs_idx ] >> 4) + (((data_a[ib].qh[qh_idx ] & hm2) != 0) ? 16 : 0)) - m2); + data_b[y_idx + 33] = D_TYPE(d2 * FLOAT_TYPE((data_a[ib].qs[qs_idx + 1] >> 4) + (((data_a[ib].qh[qh_idx + 1] & hm2) != 0) ? 16 : 0)) - m2); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp new file mode 100644 index 000000000..0b9131755 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp @@ -0,0 +1,33 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { + const uint i = gl_WorkGroupID.x * 256 + wgy; + if (i >= p.M * p.K / QUANT_K) { + return; + } + const uint tid = gl_LocalInvocationID.x; + const uint ip = tid / 32; + const uint il = tid - 32 * ip; + const uint is = 8 * ip + il / 16; + + const uint y_idx = i * QUANT_K + 128 * ip + il; + + const uint ql_idx = 64 * ip + il; + const uint8_t qh = data_a[i].qh[32 * ip + il]; + + const FLOAT_TYPE d = FLOAT_TYPE(data_a[i].d); + + data_b[y_idx + 0] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 0] * (int8_t((data_a[i].ql[ql_idx + 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32))); + data_b[y_idx + 32] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 2] * (int8_t((data_a[i].ql[ql_idx + 32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32))); + data_b[y_idx + 64] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 4] * (int8_t((data_a[i].ql[ql_idx + 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32))); + data_b[y_idx + 96] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 6] * (int8_t((data_a[i].ql[ql_idx + 32] >> 4) | (((qh >> 6) & 3) << 4)) - 32))); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp new file mode 100644 index 000000000..bd1344a88 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp @@ -0,0 +1,31 @@ +#version 450 + +#include "dequant_head.comp" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_q8_0 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint b_idx = 1024*i + 32*ir + 16*il; + + const float d = float(data_a[ib].d); + + const uint q_idx = 16*il; + + [[unroll]] for (uint l = 0; l < 16; l += 2) { + data_b[b_idx + l ] = D_TYPE(d * data_a[ib].qs[q_idx + l ]); + data_b[b_idx + l + 1] = D_TYPE(d * data_a[ib].qs[q_idx + l + 1]); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp new file mode 100644 index 000000000..4e68742b5 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp @@ -0,0 +1,34 @@ +#version 450 + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_control_flow_attributes : enable + +layout (push_constant) uniform parameter +{ + uint ncols; + uint rows_per_channel; + uint n_past; +} p; + +#include "types.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint col = gl_GlobalInvocationID.y; + const uint row = gl_GlobalInvocationID.x; + + if (col >= p.ncols) { + return; + } + + const uint i = row*p.ncols + col; + if (col > p.n_past + row % p.rows_per_channel) { + data_d[i] = D_TYPE(uintBitsToFloat(0xFF800000)); + } else { + data_d[i] = D_TYPE(data_a[i]); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp new file mode 100644 index 000000000..9fb69c6c1 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp @@ -0,0 +1,27 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" + +const uint num_threads = 256; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + uint idx = get_idx(); + + // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation + const uint num_iter = 2; + + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (idx >= p.ne) { + continue; + } + uint i00, i01, i02, i03; + get_indices(idx, i00, i01, i02, i03); + + data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) / FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); + + idx += num_threads; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp new file mode 100644 index 000000000..c5be8131b --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp @@ -0,0 +1,289 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require + +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +#extension GL_KHR_memory_scope_semantics : enable +#extension GL_KHR_cooperative_matrix : enable +#extension GL_NV_cooperative_matrix2 : enable +#extension GL_EXT_buffer_reference : enable +#extension GL_KHR_shader_subgroup_ballot : enable +#extension GL_KHR_shader_subgroup_vote : enable +#extension GL_EXT_null_initializer : enable + +#include "types.comp" +#include "dequant_funcs_cm2.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (constant_id = 1) const uint32_t Br = 32; +layout (constant_id = 2) const uint32_t Bc = 32; +layout (constant_id = 3) const uint32_t D = 32; +layout (constant_id = 4) const uint32_t Clamp = gl_CooperativeMatrixClampModeConstantNV; + +layout (push_constant) uniform parameter { + uint32_t N; + uint32_t KV; + + uint32_t ne1; + uint32_t ne2; + uint32_t ne3; + + uint32_t neq2; + uint32_t neq3; + uint32_t nek2; + uint32_t nek3; + uint32_t nev2; + uint32_t nev3; + uint32_t nem1; + + uint32_t nb02; + uint32_t nb03; + uint32_t nb12; + uint32_t nb13; + uint32_t nb22; + uint32_t nb23; + uint32_t nb31; + + float scale; + float max_bias; + float logit_softcap; + + uint32_t mask; + uint32_t n_head_log2; + float m0; + float m1; +} p; + +layout (binding = 0) readonly buffer Q {uint8_t data_q[];}; +layout (binding = 1) readonly buffer K {uint8_t data_k[];}; +layout (binding = 2) readonly buffer V {uint8_t data_v[];}; +layout (binding = 3) readonly buffer M {uint8_t data_m[];}; +layout (binding = 4) writeonly buffer O {D_TYPE data_o[];}; + +#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b)) + +ACC_TYPE maxReduce(const in ACC_TYPE x, const in ACC_TYPE y) { + return max(x, y); +} + +ACC_TYPE smearReduce(const in ACC_TYPE x, const in ACC_TYPE y) { + return x; +} + +// Replace matrix elements >= numRows or numCols with 'replace' +ACC_TYPE replacePadding(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem, const in ACC_TYPE replace, const in uint32_t numRows, const in uint32_t numCols) { + if (row >= numRows || col >= numCols) { + return replace; + } + return elem; +} + +ACC_TYPE Exp(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem) +{ + return exp(elem); +} + +ACC_TYPE Max(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem0, const in ACC_TYPE elem1) +{ + return max(elem0, elem1); +} + +#if defined(BLOCK_SIZE) +#define DECODEFUNC , DEQUANTFUNC +#else +#define DECODEFUNC +#endif + +void main() { +#if defined(DATA_A_IQ4_NL) + init_iq4nl_shmem(); +#endif + + const uint32_t N = p.N; + const uint32_t KV = p.KV; + + const uint32_t Tr = CEIL_DIV(N, Br); + const uint32_t Tc = CEIL_DIV(KV, Bc); + + const uint32_t i = gl_WorkGroupID.x; + + const uint32_t iq2 = gl_WorkGroupID.y; + const uint32_t iq3 = gl_WorkGroupID.z; + + // broadcast factors + const uint32_t rk2 = p.neq2/p.nek2; + const uint32_t rk3 = p.neq3/p.nek3; + + const uint32_t rv2 = p.neq2/p.nev2; + const uint32_t rv3 = p.neq3/p.nev3; + + // k indices + const uint32_t ik3 = iq3 / rk3; + const uint32_t ik2 = iq2 / rk2; + + // v indices + const uint32_t iv3 = iq3 / rv3; + const uint32_t iv2 = iq2 / rv2; + + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutQ = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutNV<2, Clamp> tensorLayoutK = createTensorLayoutNV(2, Clamp); + tensorLayoutNV<2, Clamp> tensorLayoutV = createTensorLayoutNV(2, Clamp); + + tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0); + +#if defined(BLOCK_SIZE) + tensorLayoutK = setTensorLayoutBlockSizeNV(tensorLayoutK, 1, BLOCK_SIZE); + tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, BLOCK_SIZE); +#endif + + tensorLayoutQ = setTensorLayoutDimensionNV(tensorLayoutQ, N, D); + tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, KV, D); + tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, KV, D); + + coopmat Q; + coopmat Qf16; + + uint32_t q_offset = iq2*p.nb02+iq3*p.nb03; + coopMatLoadTensorNV(Q, data_q, q_offset, sliceTensorLayoutNV(tensorLayoutQ, i * Br, Br, 0, D)); + + Qf16 = coopmat(Q); + Qf16 *= float16_t(p.scale); + + coopmat O = coopmat(0); + + coopmat L, M; + + L = coopmat(0); + M = coopmat(-1.0/0.0); + + ACC_TYPE slope = ACC_TYPE(1.0); + + // ALiBi + if (p.max_bias > 0.0f) { + const uint32_t h = iq2; + + const ACC_TYPE base = ACC_TYPE(h < p.n_head_log2 ? p.m0 : p.m1); + const int exph = int(h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1); + + slope = pow(base, ACC_TYPE(exph)); + } + + [[dont_unroll]] + for (uint32_t j = 0; j < Tc; ++j) { + + coopmat S = coopmat(0); + + coopmat K_T; + + uint32_t k_offset = ik2*p.nb12 + ik3*p.nb13; + coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, D), tensorViewTranspose DECODEFUNC); + S = coopMatMulAdd(Qf16, K_T, S); + + if (p.logit_softcap != 0.0f) { + [[unroll]] + for (int k = 0; k < S.length(); ++k) { + S[k] = ACC_TYPE(p.logit_softcap)*tanh(S[k]); + } + } + + if (p.mask != 0) { + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutM = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutM = setTensorLayoutDimensionNV(tensorLayoutM, p.nem1, KV); + + coopmat mv; + + coopMatLoadTensorNV(mv, data_m, 0, sliceTensorLayoutNV(tensorLayoutM, i * Br, Br, j * Bc, Bc)); + + S += slope*coopmat(mv); + } + + // Clear padding elements to -inf, so they don't contribute to rowmax + if (Clamp != 0 && + ((j + 1) * Bc > KV || + (i + 1) * Br > N)) { + + uint R = ((i + 1) * Br > N) ? (N % Br) : Br; + uint C = ((j + 1) * Bc > KV) ? (KV % Bc) : Bc; + + coopMatPerElementNV(S, S, replacePadding, ACC_TYPE(-1.0/0.0), R, C); + } + + coopmat rowmax, P, rowsum, eM; + + coopMatReduceNV(rowmax, S, gl_CooperativeMatrixReduceRowNV, maxReduce); + + coopmat Mold = M; + + // M = max(rowmax, Mold) + // P = e^(S - M) + // eM = e^(Mold - M) + coopMatPerElementNV(M, rowmax, Max, Mold); + coopMatPerElementNV(P, S - M, Exp); + coopMatPerElementNV(eM, Mold - M, Exp); + + // Clear padding elements to 0, so they don't contribute to rowsum + if (Clamp != 0 && + ((j + 1) * Bc > KV || + (i + 1) * Br > N)) { + + uint R = ((i + 1) * Br > N) ? (N % Br) : Br; + uint C = ((j + 1) * Bc > KV) ? (KV % Bc) : Bc; + + coopMatPerElementNV(P, P, replacePadding, ACC_TYPE(0.0), R, C); + } + + coopmat P_A = coopmat(P); + + // compute rowsum by multiplying by matrix of all ones. + coopmat One = coopmat(1.0); + + rowsum = coopmat(0.0); + rowsum = coopMatMulAdd(P_A, One, rowsum); + + coopmat V; + uint32_t v_offset = iv2*p.nb22 + iv3*p.nb23; + coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, D) DECODEFUNC); + + L = eM*L + rowsum; + + // This is the "diagonal" matrix in the paper, but since we do componentwise + // multiply rather than matrix multiply it has the diagonal element smeared + // across the row + coopmat eMdiag; + + // resize eM by using smear/reduce + coopMatReduceNV(eMdiag, eM, gl_CooperativeMatrixReduceRowNV, smearReduce); + + O = eMdiag * O; + + O = coopMatMulAdd(P_A, V, O); + } + + coopmat Ldiag; + + // resize L by using smear/reduce + coopMatReduceNV(Ldiag, L, gl_CooperativeMatrixReduceRowNV, smearReduce); + + [[unroll]] + for (int k = 0; k < Ldiag.length(); ++k) { + Ldiag[k] = ACC_TYPE(1.0) / Ldiag[k]; + } + + O = Ldiag*O; + + tensorLayoutNV<3, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(3, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.ne2, p.ne1, D); + + // permute dimensions + tensorViewNV<3, false, 1, 0, 2> tensorViewPermute = createTensorViewNV(3, false, 1, 0, 2); + uint32_t o_offset = iq3*p.ne2*p.ne1; + + coopmat O_D = coopmat(O); + coopMatStoreTensorNV(O_D, data_o, o_offset, sliceTensorLayoutNV(tensorLayoutD, i * Br, Br, iq2, 1, 0, D), tensorViewPermute); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp new file mode 100644 index 000000000..4cc7a68ca --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp @@ -0,0 +1,25 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const float GELU_COEF_A = 0.044715f; + const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + + const float xi = float(data_a[i]); + const float val = SQRT_2_OVER_PI*xi*(1.0f + GELU_COEF_A*xi*xi); + data_d[i] = D_TYPE(0.5f*xi*(2.0f - 2.0f / (exp(2 * val) + 1))); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp new file mode 100644 index 000000000..e6e6fcfd2 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp @@ -0,0 +1,23 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const float GELU_QUICK_COEF = -1.702f; + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + + const float x = float(data_a[i]); + data_d[i] = D_TYPE(x * (1.0f / (1.0f + exp(GELU_QUICK_COEF * x)))); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp new file mode 100644 index 000000000..062e2a4cd --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp @@ -0,0 +1,64 @@ +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_control_flow_attributes : require + +layout (push_constant) uniform parameter +{ + uint ne; + uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03; + uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13; + uint ne20; uint ne21; uint ne22; uint ne23; uint nb20; uint nb21; uint nb22; uint nb23; + uint misalign_offsets; + float param1; float param2; int param3; +} p; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; + +// true if src0/src1 are the same shape and the indices can be reused without additional modulus +layout(constant_id = 0) const bool norepeat = false; + +uint get_idx() { + return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; +} + +uint get_aoffset() { return p.misalign_offsets >> 16; } +uint get_boffset() { return (p.misalign_offsets >> 8) & 0xFF; } +uint get_doffset() { return p.misalign_offsets & 0xFF; } + +// mod and div are expensive and coordinates/dimensions are often power of 2 or equal to 1 +uint fastmod(uint a, uint b) { + if ((b & (b-1)) == 0) { + return a & (b-1); + } + return a % b; +} + +uint fastdiv(uint a, uint b) { + return (a < b) ? 0 : (a / b); +} + +void get_indices(uint idx, out uint i00, out uint i01, out uint i02, out uint i03) { + i03 = fastdiv(idx, (p.ne02*p.ne01*p.ne00)); + const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00; + i02 = fastdiv((idx - i03_offset), (p.ne01*p.ne00)); + const uint i02_offset = i02*p.ne01*p.ne00; + i01 = (idx - i03_offset - i02_offset) / p.ne00; + i00 = idx - i03_offset - i02_offset - i01*p.ne00; +} + +uint src0_idx(uint i00, uint i01, uint i02, uint i03) { + return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00; +} + +uint src1_idx(uint i00, uint i01, uint i02, uint i03) { + if (norepeat) { + return i03*p.nb13 + i02*p.nb12 + i01*p.nb11 + i00*p.nb10; + } else { + return fastmod(i03, p.ne13)*p.nb13 + fastmod(i02, p.ne12)*p.nb12 + fastmod(i01, p.ne11)*p.nb11 + fastmod(i00, p.ne10)*p.nb10; + } +} + +uint dst_idx(uint i00, uint i01, uint i02, uint i03) { + return i03*p.nb23 + i02*p.nb22 + i01*p.nb21 + i00*p.nb20; +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp new file mode 100644 index 000000000..66e46ae67 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp @@ -0,0 +1,9 @@ +#extension GL_EXT_shader_16bit_storage : require + +layout (push_constant) uniform parameter +{ + uint KX; + uint KY; + float param1; + float param2; +} p; diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp new file mode 100644 index 000000000..68d1bc9f1 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp @@ -0,0 +1,56 @@ +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_control_flow_attributes : require + +layout (push_constant) uniform parameter +{ + uint ne; + uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03; + uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13; + uint misalign_offsets; + float param1; float param2; + + uint ne0_012mp; uint ne0_012L; + uint ne0_01mp; uint ne0_01L; + uint ne0_0mp; uint ne0_0L; + uint ne1_012mp; uint ne1_012L; + uint ne1_01mp; uint ne1_01L; + uint ne1_0mp; uint ne1_0L; +} p; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +uint get_idx() { + return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; +} + +uint get_aoffset() { return p.misalign_offsets >> 16; } +uint get_doffset() { return p.misalign_offsets & 0xFFFF; } + +// see init_fastdiv_values in ggml-vulkan.cpp +uint fastdiv(uint n, uint mp, uint L) { + uint msbs, lsbs; + // msbs = mulhi(n, mp) + umulExtended(n, mp, msbs, lsbs); + return (msbs + n) >> L; +} + +uint src0_idx(uint idx) { + const uint i03 = fastdiv(idx, p.ne0_012mp, p.ne0_012L); + const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00; + const uint i02 = fastdiv(idx - i03_offset, p.ne0_01mp, p.ne0_01L); + const uint i02_offset = i02*p.ne01*p.ne00; + const uint i01 = fastdiv(idx - i03_offset - i02_offset, p.ne0_0mp, p.ne0_0L); + const uint i00 = idx - i03_offset - i02_offset - i01*p.ne00; + return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00; +} + +uint dst_idx(uint idx) { + const uint i13 = fastdiv(idx, p.ne1_012mp, p.ne1_012L); + const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10; + const uint i12 = fastdiv(idx - i13_offset, p.ne1_01mp, p.ne1_01L); + const uint i12_offset = i12*p.ne11*p.ne10; + const uint i11 = fastdiv(idx - i13_offset - i12_offset, p.ne1_0mp, p.ne1_0L); + const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10; + return i13*p.nb13 + i12*p.nb12 + i11*p.nb11 + i10*p.nb10; +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp new file mode 100644 index 000000000..e877ed779 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp @@ -0,0 +1,28 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint i00 = gl_GlobalInvocationID.x; + const uint i10 = gl_GlobalInvocationID.y; + const uint i11 = (gl_GlobalInvocationID.z)/p.ne12; + const uint i12 = (gl_GlobalInvocationID.z)%p.ne12; + + if (i00 >= p.ne00) { + return; + } + + const uint i01 = data_b[get_boffset() + i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; + + const uint a_offset = get_aoffset() + i01*p.nb01 + i11*p.nb02 + i12*p.nb03; + const uint d_offset = get_doffset() + i10*p.nb21 + i11*p.nb22 + i12*p.nb23; + +#ifndef OPTIMIZATION_ERROR_WORKAROUND + data_d[d_offset + i00] = D_TYPE(data_a[a_offset + i00]); +#else + data_d[d_offset + i00] = data_a[a_offset + i00]; +#endif +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp new file mode 100644 index 000000000..1426fde65 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp @@ -0,0 +1,39 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" +#include "dequant_funcs.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint i00 = (gl_GlobalInvocationID.x)*2; + const uint i10 = gl_GlobalInvocationID.y; + const uint i11 = (gl_GlobalInvocationID.z)/p.ne12; + const uint i12 = (gl_GlobalInvocationID.z)%p.ne12; + +#if defined(DATA_A_IQ4_NL) + init_iq4nl_shmem(); +#endif + + if (i00 >= p.ne00) { + return; + } + + const uint i01 = data_b[i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; + + const uint a_offset = i01*p.nb01 + i11*p.nb02 + i12*p.nb03; + const uint d_offset = i10*p.nb21 + i11*p.nb22 + i12*p.nb23; + + const uint ib = a_offset + i00/QUANT_K; // block index + const uint iqs = (i00%QUANT_K)/QUANT_R; // quant index + const uint iybs = i00 - i00%QUANT_K; // dst block start index + const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2; + + vec2 v = dequantize(ib, iqs, 0); + const vec2 dm = get_dm(ib, 0); + v = v * dm.x + dm.y; + + data_d[d_offset + iybs + iqs ] = D_TYPE(v.x); + data_d[d_offset + iybs + iqs + y_offset] = D_TYPE(v.y); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp new file mode 100644 index 000000000..b6a0d5645 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp @@ -0,0 +1,66 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable +#define BLOCK_SIZE 512 + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +shared float tmp[BLOCK_SIZE]; + +void main() { + const uint group_size = p.KX; + const float eps = p.param1; + + const uint tid = gl_LocalInvocationID.x; + const uint start = gl_WorkGroupID.x * group_size + tid; + const uint end = (gl_WorkGroupID.x + 1) * group_size; + + tmp[tid] = 0.0f; + + // Calculate mean + [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { + tmp[tid] += float(data_a[col]); + } + + // tmp up partial tmps and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] += tmp[tid + s]; + } + barrier(); + } + + const float mean = tmp[0] / group_size; + barrier(); + tmp[tid] = 0.0f; + + // Calculate variance + [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { + const float xi = float(data_a[col]) - mean; + data_d[col] = D_TYPE(xi); + tmp[tid] += xi * xi; + } + + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] += tmp[tid + s]; + } + barrier(); + } + + const float variance = tmp[0] / group_size; + const float scale = inversesqrt(variance + eps); + + [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { + data_d[col] *= D_TYPE(scale); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp new file mode 100644 index 000000000..122b1e93f --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp @@ -0,0 +1,87 @@ +#version 450 + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_spirv_intrinsics: enable +#extension GL_EXT_control_flow_attributes : require + +#if RTE16 +spirv_execution_mode(capabilities = [4467], 4462, 16); // RoundingModeRTE, 16 bits +#endif + +layout (push_constant) uniform parameter +{ + uint batch_offset; uint offset_delta; + uint IC; + uint IW; uint IH; + uint OW; uint OH; + uint KW; uint KH; + uint pelements; + uint CHW; + int s0; int s1; + int p0; int p1; + int d0; int d1; +} p; + +#include "types.comp" + +layout(constant_id = 0) const uint BLOCK_SIZE = 32; + +const uint NUM_ITER = 512 / BLOCK_SIZE; + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint gidx = gl_GlobalInvocationID.x; + + const uint oh = gl_GlobalInvocationID.y; + const uint batch = gl_GlobalInvocationID.z / p.IC; + const uint ic = gl_GlobalInvocationID.z % p.IC; + + A_TYPE values[NUM_ITER]; + uint offset_dst[NUM_ITER]; + [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { + values[idx] = A_TYPE(0); + } + + [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { + + const uint i = gidx * NUM_ITER + idx; + + const uint ksize = p.OW * (p.KH > 1 ? p.KW : 1); + const uint kx = i / ksize; + const uint kd = kx * ksize; + const uint ky = (i - kd) / p.OW; + const uint ix = i % p.OW; + + const uint iiw = ix * p.s0 + kx * p.d0 - p.p0; + const uint iih = oh * p.s1 + ky * p.d1 - p.p1; + + offset_dst[idx] = + ((batch * p.OH + oh) * p.OW + ix) * p.CHW + + (ic * (p.KW * p.KH) + ky * p.KW + kx); + + if (i >= p.pelements) { + continue; + } + + if (iih < p.IH && iiw < p.IW) { + const uint offset_src = ic * p.offset_delta + batch * p.batch_offset; + values[idx] = data_a[offset_src + iih * p.IW + iiw]; + } + } + + [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { + + const uint i = gidx * NUM_ITER + idx; + + if (i >= p.pelements) { + continue; + } + + data_d[offset_dst[idx]] = D_TYPE(values[idx]); + } + +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp new file mode 100644 index 000000000..d90a99aea --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp @@ -0,0 +1,22 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + + const float val = float(data_a[i]); + data_d[i] = D_TYPE(max(val, 0.0f) + min(val, 0.0f) * p.param1); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp new file mode 100644 index 000000000..43de19df8 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp @@ -0,0 +1,27 @@ +#version 450 + +#include "types.comp" +#include "generic_binary_head.comp" + +const uint num_threads = 256; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + uint idx = get_idx(); + + // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation + const uint num_iter = 2; + + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (idx >= p.ne) { + continue; + } + uint i00, i01, i02, i03; + get_indices(idx, i00, i01, i02, i03); + + data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) * FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); + + idx += num_threads; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp new file mode 100644 index 000000000..4c64fd47a --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp @@ -0,0 +1,48 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {float data_a[];}; +layout (binding = 0) readonly buffer A4 {vec4 data_a4[];}; +layout (binding = 1) writeonly buffer D {float data_d[];}; +layout (binding = 1) writeonly buffer D4 {vec4 data_d4[];}; + +layout (push_constant) uniform parameter { + uint ne; + uint k_num; +} p; + +void main() { + // Each invocation handles four consecutive components + const uint idx = gl_GlobalInvocationID.x * 4; + + if (idx >= p.ne) { + return; + } + + // Check if all four components are in bounds and aligned, + // then use vector loads + if (idx + 3 < p.ne && (p.ne % 4) == 0) { + vec4 result = vec4(0.0f); + + [[unroll]] for (uint i = 0; i < p.k_num; i++) { + result += data_a4[(i * p.ne + idx) / 4]; + } + + data_d4[idx / 4] = result; + } else { + [[unroll]] for (uint j = 0; j < 4; ++j) { + if (idx + j < p.ne) { + float result = 0.0f; + + [[unroll]] for (uint i = 0; i < p.k_num; i++) { + result += data_a[i * p.ne + idx + j]; + } + + data_d[idx + j] = result; + } + } + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp new file mode 100644 index 000000000..24875cdcf --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp @@ -0,0 +1,152 @@ +#version 450 + +#ifdef FLOAT16 +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#endif +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +#if !defined(DATA_A_F32) && !defined(DATA_A_F16) +#define K_PER_ITER 8 +#else +#define K_PER_ITER 2 +#endif + + +uint a_offset, b_offset, d_offset, y_offset; + +void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint tid, const uint i, bool lastiter) +{ + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + const uint col = i*BLOCK_SIZE + K_PER_ITER*tid; + const uint iqs = (col%QUANT_K)/QUANT_R; // quant index + const uint iybs = col - col%QUANT_K; // y block start index + +#if K_PER_ITER == 8 +#if QUANT_R == 2 + const B_TYPE_VEC4 bv02 = data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]; + const B_TYPE_VEC4 bv13 = data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs + y_offset) / 4]; + const vec4 bv0 = vec4(bv02.x, bv13.x, bv02.y, bv13.y); + const vec4 bv1 = vec4(bv02.z, bv13.z, bv02.w, bv13.w); +#else + const vec4 bv0 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]); + const vec4 bv1 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4 + 1]); +#endif +#else + // Check if the second of the pair of elements is OOB, and don't fetch B or + // accumulate it. We still fetch a pair of elements for A, which is fine for + // quantized formats since they'll be within the same block. We should + // probably skip fetching the second element for F16/F32, but as of now we + // still do. + const bool OOB = lastiter && (iybs + iqs + y_offset >= p.ncols); + + FLOAT_TYPE b0 = 0, b1 = 0; + b0 = FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs]); + if (!OOB) { + b1 = FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs + y_offset]); + } +#endif + uint ibi = first_row*p.ncols; + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib = (ibi + col)/QUANT_K; // block index + ibi += p.ncols; + +#if K_PER_ITER == 8 + vec4 v = dequantize4(ib, iqs, a_offset); + vec4 v2 = dequantize4(ib, iqs+(4/QUANT_R), a_offset); + + const vec2 dm = get_dm(ib, a_offset); + if (dm.y != 0) { // quant has min component + v = v * dm.x + dm.y; + v2 = v2 * dm.x + dm.y; + } + + // matrix multiplication + FLOAT_TYPE rowtmp = dot(bv0, v); + rowtmp += dot(bv1, v2); + + if (dm.y == 0) + rowtmp *= dm.x; + + temp[j][n] += rowtmp; +#else + const vec2 v = dequantize(ib, iqs, a_offset); + + // matrix multiplication + temp[j][n] = fma(FLOAT_TYPE(v.x), b0, temp[j][n]); + if (!OOB) { + temp[j][n] = fma(FLOAT_TYPE(v.y), b1, temp[j][n]); + } +#endif + } + } +} + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + const uint tid = gl_LocalInvocationID.x; + + get_offsets(a_offset, b_offset, d_offset); + a_offset /= QUANT_K; + + y_offset = QUANT_R == 1 ? 1 : QUANT_K/2; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + uint num_iters = p.ncols / (K_PER_ITER * BLOCK_SIZE); + if (num_iters * K_PER_ITER * BLOCK_SIZE + K_PER_ITER*tid < p.ncols) { + num_iters++; + } + int unroll_count = 4; + uint unrolled_iters = num_iters & ~(unroll_count - 1); + + uint i = 0; + while (i < unrolled_iters) { + // Manually partially unroll the loop + [[unroll]] for (uint k = 0; k < unroll_count; ++k) { + iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false); + i++; + } + } + unroll_count = 2; + unrolled_iters = num_iters & ~(unroll_count - 1); + while (i < unrolled_iters) { + // Manually partially unroll the loop + [[unroll]] for (uint k = 0; k < unroll_count; ++k) { + iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false); + i++; + } + } + while (i < num_iters) { + iter(temp, first_row, num_rows, tid, i*K_PER_ITER, true); + i++; + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + +#if defined(DATA_A_IQ4_NL) + init_iq4nl_shmem(); +#endif + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp new file mode 100644 index 000000000..903753c7e --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp @@ -0,0 +1,118 @@ +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_8bit_storage : require + +#ifdef MUL_MAT_ID +#define EXPERT_COUNT 8 +#endif + +#include "types.comp" + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 1) readonly buffer BV2 {B_TYPE_VEC2 data_b_v2[];}; +layout (binding = 1) readonly buffer BV4 {B_TYPE_VEC4 data_b_v4[];}; + +layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; +#ifdef MUL_MAT_ID +layout (binding = 3) readonly buffer IDS {int data_ids[];}; +#endif + +#include "dequant_funcs.comp" + +layout (push_constant) uniform parameter +{ + uint ncols; + uint stride_a; + uint stride_b; + uint stride_d; + + uint batch_stride_a; + uint batch_stride_b; + uint batch_stride_d; + +#ifdef MUL_MAT_ID + uint nei0; + uint ne11; +#else + uint ne02; + uint ne12; + uint broadcast2; + uint broadcast3; +#endif +} p; + +void get_offsets(out uint a_offset, out uint b_offset, out uint d_offset) { +#ifdef MUL_MAT_ID + const uint expert_idx = gl_GlobalInvocationID.y; +#else + const uint batch_idx = gl_GlobalInvocationID.y; +#endif + +#ifndef MUL_MAT_ID + uint batch_idx_a = 0; + if (batch_idx != 0) { + const uint i13 = batch_idx / p.ne12; + const uint i12 = batch_idx % p.ne12; + + const uint i03 = i13 / p.broadcast3; + const uint i02 = i12 / p.broadcast2; + + batch_idx_a = i03 * p.ne02 + i02; + } +#else + const uint expert_id = data_ids[expert_idx]; +#endif + + a_offset = +#ifdef MUL_MAT_ID + expert_id * p.batch_stride_a; +#else + batch_idx_a * p.batch_stride_a; +#endif + b_offset = +#ifdef MUL_MAT_ID + (expert_idx % p.ne11) * p.stride_b; +#else + batch_idx * p.batch_stride_b; +#endif + d_offset = +#ifdef MUL_MAT_ID + expert_idx * p.stride_d; +#else + batch_idx * p.batch_stride_d; +#endif +} + +layout (constant_id = 0) const uint BLOCK_SIZE = 32; +layout (constant_id = 1) const uint NUM_ROWS = 1; +layout (constant_id = 2) const uint NUM_COLS = 1; + +shared FLOAT_TYPE tmpsh[NUM_COLS][NUM_ROWS][BLOCK_SIZE]; + +void reduce_result(const in FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const in uint32_t d_offset, const in uint32_t first_row, const in uint32_t num_rows, const in uint32_t tid) { + // sum up partial sums and write back result + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + tmpsh[j][n][tid] = temp[j][n]; + } + } + barrier(); + [[unroll]] for (uint s = BLOCK_SIZE/2; s > 0; s >>= 1) { + if (tid < s) { + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + tmpsh[j][n][tid] += tmpsh[j][n][tid + s]; + } + } + } + barrier(); + } + if (tid == 0) { + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + data_d[j*p.batch_stride_d + d_offset + first_row + n] = D_TYPE(tmpsh[j][n][0]); + } + } + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp new file mode 100644 index 000000000..1cc4996d3 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp @@ -0,0 +1,71 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require + +#define BLOCK_SIZE 32 +#define FLOAT_TYPE float + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE dst[];}; + +layout (push_constant) uniform parameter +{ + uint ncols_x; + uint nrows_x; + uint row_stride_x; + uint channel_stride_x; + uint channel_x_divisor; + uint b_offset; + uint d_offset; +} p; + +shared FLOAT_TYPE tmp[BLOCK_SIZE]; + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint row_x = gl_GlobalInvocationID.y; + const uint channel = gl_GlobalInvocationID.z; + const uint channel_x = channel / p.channel_x_divisor; + + const uint nrows_y = p.ncols_x; + const uint nrows_dst = p.nrows_x; + const uint row_dst = row_x; + + const uint idst = channel*nrows_dst + row_dst; + + tmp[tid] = 0.0f; + + for (uint col_x0 = 0; col_x0 < p.ncols_x; col_x0 += BLOCK_SIZE) { + const uint col_x = col_x0 + tid; + + if (col_x >= p.ncols_x) { + break; + } + + const uint row_y = col_x; + + const uint ix = channel_x*p.channel_stride_x + row_x*p.row_stride_x + col_x; + const uint iy = channel*nrows_y + row_y; + + const FLOAT_TYPE xi = FLOAT_TYPE(data_a[ix]); + + tmp[tid] = fma(xi, FLOAT_TYPE(data_b[iy]), tmp[tid]); + } + + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] += tmp[tid + s]; + } + barrier(); + } + + if (tid == 0) { + dst[idst] = tmp[0]; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp new file mode 100644 index 000000000..9b443807d --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp @@ -0,0 +1,73 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require + +#define BLOCK_SIZE 32 +#define FLOAT_TYPE float + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE dst[];}; + +layout (push_constant) uniform parameter +{ + uint ncols_x; + uint nrows_x; + uint nchannels_x; + uint nchannels_y; + uint b_offset; + uint d_offset; +} p; + +shared FLOAT_TYPE tmp[BLOCK_SIZE]; + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint row_x = gl_GlobalInvocationID.y; + const uint channel = gl_GlobalInvocationID.z; + const uint channel_x = channel / (p.nchannels_y / p.nchannels_x); + + const uint nrows_y = p.ncols_x; + const uint nrows_dst = p.nrows_x; + const uint row_dst = row_x; + + tmp[tid] = FLOAT_TYPE(0.0f); + + for (uint col_x0 = 0; col_x0 < p.ncols_x; col_x0 += BLOCK_SIZE) { + const uint col_x = col_x0 + tid; + + if (col_x >= p.ncols_x) { + break; + } + + // x is transposed and permuted + const uint ix = row_x*p.nchannels_x*p.ncols_x + channel_x*p.ncols_x + col_x; + const FLOAT_TYPE xi = FLOAT_TYPE(data_a[ix]); + + const uint row_y = col_x; + + // y is not transposed but permuted + const uint iy = channel*nrows_y + row_y; + + tmp[tid] = fma(xi, FLOAT_TYPE(data_b[iy]), tmp[tid]); + } + + // dst is not transposed and not permuted + const uint idst = channel*nrows_dst + row_dst; + + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] += tmp[tid + s]; + } + barrier(); + } + + if (tid == 0) { + dst[idst] = tmp[0]; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp new file mode 100644 index 000000000..934213446 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp @@ -0,0 +1,115 @@ +#version 450 +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads are used to process each block + const uint it_size = gl_WorkGroupSize.x/16; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid%16; // 0...16 + const uint ix = tid/16; + + const uint step = 8; + + const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... + const uint v_in = itid - step*v_im; // 0...15 or 0...7 + + const uint l0 = 2*v_in; // 0...15 + const uint q_offset = 32*v_im + l0; + const uint s_offset = 8*v_im; + const uint y_offset = 128*v_im + l0; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { + const uint y_idx = i * QUANT_K + y_offset; + + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; + f16vec2 d = data_a[ib0 + i].d; + const FLOAT_TYPE dall = d.x; + const FLOAT_TYPE dmin = d.y; + + uint32_t s0_u32 = data_a_packed32[ib0 + i].scales[s_offset / 4 + 0]; + uint32_t s4_u32 = data_a_packed32[ib0 + i].scales[s_offset / 4 + 1]; + + uint32_t s0_lo4_u32 = s0_u32 & 0x0F0F0F0F; + uint32_t s0_hi4_u32 = (s0_u32 >> 4) & 0x0F0F0F0F; + uint32_t s4_lo4_u32 = s4_u32 & 0x0F0F0F0F; + uint32_t s4_hi4_u32 = (s4_u32 >> 4) & 0x0F0F0F0F; + + uvec4 s0_lo4 = uvec4(unpack8(s0_lo4_u32)); + uvec4 s4_lo4 = uvec4(unpack8(s4_lo4_u32)); + uvec4 s0_hi4 = uvec4(unpack8(s0_hi4_u32)); + uvec4 s4_hi4 = uvec4(unpack8(s4_hi4_u32)); + + uint16_t qs0_u16 = data_a_packed16[ib0 + i].qs[q_offset / 2 + 0]; + uint16_t qs16_u16 = data_a_packed16[ib0 + i].qs[q_offset / 2 + 8]; + uvec2 qs0 = uvec2(unpack8(qs0_u16)); + uvec2 qs16 = uvec2(unpack8(qs16_u16)); + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + B_TYPE_VEC2 b0 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 0]; + B_TYPE_VEC2 b16 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 8]; + B_TYPE_VEC2 b32 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 16]; + B_TYPE_VEC2 b48 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 24]; + B_TYPE_VEC2 b64 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 32]; + B_TYPE_VEC2 b80 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 40]; + B_TYPE_VEC2 b96 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 48]; + B_TYPE_VEC2 b112 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 56]; + + FLOAT_TYPE sum1 = FLOAT_TYPE(0.0); + FLOAT_TYPE sum2 = FLOAT_TYPE(0.0); + [[unroll]] for (int l = 0; l < 2; ++l) { + sum1 = fma(FLOAT_TYPE(b0[l]), FLOAT_TYPE(s0_lo4[0]) * FLOAT_TYPE((qs0[l] >> 0) & 3), + fma(FLOAT_TYPE(b16[l]), FLOAT_TYPE(s0_lo4[1]) * FLOAT_TYPE((qs16[l] >> 0) & 3), + fma(FLOAT_TYPE(b32[l]), FLOAT_TYPE(s0_lo4[2]) * FLOAT_TYPE((qs0[l] >> 2) & 3), + fma(FLOAT_TYPE(b48[l]), FLOAT_TYPE(s0_lo4[3]) * FLOAT_TYPE((qs16[l] >> 2) & 3), + fma(FLOAT_TYPE(b64[l]), FLOAT_TYPE(s4_lo4[0]) * FLOAT_TYPE((qs0[l] >> 4) & 3), + fma(FLOAT_TYPE(b80[l]), FLOAT_TYPE(s4_lo4[1]) * FLOAT_TYPE((qs16[l] >> 4) & 3), + fma(FLOAT_TYPE(b96[l]), FLOAT_TYPE(s4_lo4[2]) * FLOAT_TYPE((qs0[l] >> 6) & 3), + fma(FLOAT_TYPE(b112[l]), FLOAT_TYPE(s4_lo4[3]) * FLOAT_TYPE((qs16[l] >> 6) & 3), sum1)))))))); + sum2 = fma(FLOAT_TYPE(b0[l]), FLOAT_TYPE(s0_hi4[0]), + fma(FLOAT_TYPE(b16[l]), FLOAT_TYPE(s0_hi4[1]), + fma(FLOAT_TYPE(b32[l]), FLOAT_TYPE(s0_hi4[2]), + fma(FLOAT_TYPE(b48[l]), FLOAT_TYPE(s0_hi4[3]), + fma(FLOAT_TYPE(b64[l]), FLOAT_TYPE(s4_hi4[0]), + fma(FLOAT_TYPE(b80[l]), FLOAT_TYPE(s4_hi4[1]), + fma(FLOAT_TYPE(b96[l]), FLOAT_TYPE(s4_hi4[2]), + fma(FLOAT_TYPE(b112[l]), FLOAT_TYPE(s4_hi4[3]), sum2)))))))); + } + temp[j][n] = fma(dall, sum1, fma(-dmin, sum2, temp[j][n])); + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp new file mode 100644 index 000000000..86b0159d9 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp @@ -0,0 +1,103 @@ +#version 450 +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads are used to process each block + const uint it_size = gl_WorkGroupSize.x/16; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid%16; // 0...16 + const uint ix = tid/16; + + const uint step = 8; + + const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... + const uint v_in = itid - step*v_im; // 0...15 or 0...7 + + const uint8_t m = uint8_t(1 << (4 * v_im)); + + const uint l0 = 2*v_in; // 0...15 + const uint q_offset = 32*v_im + l0; + const uint y_offset = 128*v_im + l0; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + const uint s_shift = 4 * v_im; + + [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { + const uint y_idx = i * QUANT_K + y_offset; + + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; + const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d); + + uint16_t s0_16 = data_a_packed16[ib0 + i].scales[0]; + uint16_t s2_16 = data_a_packed16[ib0 + i].scales[1]; + uint16_t s4_16 = data_a_packed16[ib0 + i].scales[2]; + uint16_t s6_16 = data_a_packed16[ib0 + i].scales[3]; + uint16_t s8_16 = data_a_packed16[ib0 + i].scales[4]; + uint16_t s10_16 = data_a_packed16[ib0 + i].scales[5]; + u8vec2 s0 = unpack8(s0_16); + u8vec2 s2 = unpack8(s2_16); + u8vec2 s4 = unpack8(s4_16); + u8vec2 s6 = unpack8(s6_16); + u8vec2 s8 = unpack8(s8_16); + u8vec2 s10 = unpack8(s10_16); + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + + B_TYPE_VEC2 b0 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 0]; + B_TYPE_VEC2 b16 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 8]; + B_TYPE_VEC2 b32 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 16]; + B_TYPE_VEC2 b48 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 24]; + B_TYPE_VEC2 b64 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 32]; + B_TYPE_VEC2 b80 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 40]; + B_TYPE_VEC2 b96 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 48]; + B_TYPE_VEC2 b112 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 56]; + + FLOAT_TYPE sum = FLOAT_TYPE(0.0); + [[unroll]] for (int l = 0; l < 2; ++l) { + sum = fma(FLOAT_TYPE(b0[l]) * FLOAT_TYPE(int8_t(((s0[0] >> s_shift) & 0xF) | ((s8[0] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 0)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b32[l]) * FLOAT_TYPE(int8_t(((s2[0] >> s_shift) & 0xF) | ((s10[0] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 1)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b64[l]) * FLOAT_TYPE(int8_t(((s4[0] >> s_shift) & 0xF) | ((s8[0] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 2)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b96[l]) * FLOAT_TYPE(int8_t(((s6[0] >> s_shift) & 0xF) | ((s10[0] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 3)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b16[l]) * FLOAT_TYPE(int8_t(((s0[1] >> s_shift) & 0xF) | ((s8[1] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 0)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b48[l]) * FLOAT_TYPE(int8_t(((s2[1] >> s_shift) & 0xF) | ((s10[1] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 1)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b80[l]) * FLOAT_TYPE(int8_t(((s4[1] >> s_shift) & 0xF) | ((s8[1] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 2)) != 0) ? 0 : 4)), + fma(FLOAT_TYPE(b112[l]) * FLOAT_TYPE(int8_t(((s6[1] >> s_shift) & 0xF) | ((s10[1] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 3)) != 0) ? 0 : 4)), sum)))))))); + } + temp[j][n] = fma(d, sum, temp[j][n]); + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp new file mode 100644 index 000000000..cd1dd8e89 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp @@ -0,0 +1,133 @@ +#version 450 + +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads are used to process each block + const uint it_size = gl_WorkGroupSize.x/16; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid%16; // 0...16 + const uint ix = tid/16; + + const uint step = 4; + + const uint il = itid/step; // 0...3 + const uint ir = itid - step*il; // 0...7 or 0...3 + const uint n = 4; + + const uint v_im = il / 2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224 + const uint v_in = il % 2; + + const uint l0 = n * (2 * ir + v_in); // 0...15 + const uint q_offset = 32*v_im + l0; + const uint y_offset = 64*v_im + l0; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { + const uint y1_idx = i * QUANT_K + y_offset; + const uint y2_idx = y1_idx + 128; + + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; + f16vec2 d = data_a[ib0 + i].d; + const FLOAT_TYPE dall = FLOAT_TYPE(d.x); + const FLOAT_TYPE dmin = FLOAT_TYPE(d.y); + + uint32_t scale0_u32 = data_a_packed16[ib0 + i].scales[v_im ]; + uint32_t scale4_u32 = data_a_packed16[ib0 + i].scales[v_im + 2]; + uint32_t scale8_u32 = data_a_packed16[ib0 + i].scales[v_im + 4]; + uvec4 scale0 = uvec4(unpack8(scale0_u32)); + uvec4 scale4 = uvec4(unpack8(scale4_u32)); + uvec4 scale8 = uvec4(unpack8(scale8_u32)); + + const uint32_t sc0 = ( scale0.x & 0x3f); + const uint32_t sc1 = ( scale0.y & 0x3f); + const uint32_t sc2 = ( scale4.x & 0x3f); + const uint32_t sc3 = ( scale4.y & 0x3f); + const uint32_t sc4 = (( scale8.x & 0x0f) | ((scale0.x & 0xc0) >> 2)); + const uint32_t sc5 = (( scale8.y & 0x0f) | ((scale0.y & 0xc0) >> 2)); + const uint32_t sc6 = (((scale8.x >> 4) & 0x0f) | ((scale4.x & 0xc0) >> 2)); + const uint32_t sc7 = (((scale8.y >> 4) & 0x0f) | ((scale4.y & 0xc0) >> 2)); + + uint32_t qs0_u32 = data_a_packed32[ib0 + i].qs[q_offset / 4]; + uint32_t qs64_u32 = data_a_packed32[ib0 + i].qs[q_offset / 4 + 16]; + + uint32_t qs0_u32_lo4 = qs0_u32 & 0x0F0F0F0F; + uint32_t qs0_u32_hi4 = (qs0_u32 >> 4) & 0x0F0F0F0F; + uint32_t qs64_u32_lo4 = qs64_u32 & 0x0F0F0F0F; + uint32_t qs64_u32_hi4 = (qs64_u32 >> 4) & 0x0F0F0F0F; + + uvec4 qs0_lo4 = uvec4(unpack8(qs0_u32_lo4)); + uvec4 qs64_lo4 = uvec4(unpack8(qs64_u32_lo4)); + uvec4 qs0_hi4 = uvec4(unpack8(qs0_u32_hi4)); + uvec4 qs64_hi4 = uvec4(unpack8(qs64_u32_hi4)); + + const uint32_t q4_0 = qs0_lo4.x; + const uint32_t q4_1 = qs0_lo4.y; + const uint32_t q4_2 = qs0_lo4.z; + const uint32_t q4_3 = qs0_lo4.w; + const uint32_t q4_4 = qs0_hi4.x; + const uint32_t q4_5 = qs0_hi4.y; + const uint32_t q4_6 = qs0_hi4.z; + const uint32_t q4_7 = qs0_hi4.w; + const uint32_t q4_8 = qs64_lo4.x; + const uint32_t q4_9 = qs64_lo4.y; + const uint32_t q4_10 = qs64_lo4.z; + const uint32_t q4_11 = qs64_lo4.w; + const uint32_t q4_12 = qs64_hi4.x; + const uint32_t q4_13 = qs64_hi4.y; + const uint32_t q4_14 = qs64_hi4.z; + const uint32_t q4_15 = qs64_hi4.w; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + B_TYPE_VEC4 by10 = data_b_v4[(j*p.batch_stride_b + b_offset + y1_idx) / 4]; + B_TYPE_VEC4 by132 = data_b_v4[(j*p.batch_stride_b + b_offset + y1_idx) / 4 + 8]; + B_TYPE_VEC4 by20 = data_b_v4[(j*p.batch_stride_b + b_offset + y2_idx) / 4]; + B_TYPE_VEC4 by232 = data_b_v4[(j*p.batch_stride_b + b_offset + y2_idx) / 4 + 8]; + + const FLOAT_TYPE sx = fma(FLOAT_TYPE(by10.x), q4_0, fma(FLOAT_TYPE(by10.y), q4_1, fma(FLOAT_TYPE(by10.z), q4_2, FLOAT_TYPE(by10.w) * q4_3))); + const FLOAT_TYPE sy = fma(FLOAT_TYPE(by132.x), q4_4, fma(FLOAT_TYPE(by132.y), q4_5, fma(FLOAT_TYPE(by132.z), q4_6, FLOAT_TYPE(by132.w) * q4_7))); + const FLOAT_TYPE sz = fma(FLOAT_TYPE(by20.x), q4_8, fma(FLOAT_TYPE(by20.y), q4_9, fma(FLOAT_TYPE(by20.z), q4_10, FLOAT_TYPE(by20.w) * q4_11))); + const FLOAT_TYPE sw = fma(FLOAT_TYPE(by232.x), q4_12, fma(FLOAT_TYPE(by232.y), q4_13, fma(FLOAT_TYPE(by232.z), q4_14, FLOAT_TYPE(by232.w) * q4_15))); + const FLOAT_TYPE smin = + fma(FLOAT_TYPE(by10.x), sc2, fma(FLOAT_TYPE(by132.x), sc3, fma(FLOAT_TYPE(by20.x), sc6, fma(FLOAT_TYPE(by232.x), sc7, + fma(FLOAT_TYPE(by10.y), sc2, fma(FLOAT_TYPE(by132.y), sc3, fma(FLOAT_TYPE(by20.y), sc6, fma(FLOAT_TYPE(by232.y), sc7, + fma(FLOAT_TYPE(by10.z), sc2, fma(FLOAT_TYPE(by132.z), sc3, fma(FLOAT_TYPE(by20.z), sc6, fma(FLOAT_TYPE(by232.z), sc7, + fma(FLOAT_TYPE(by10.w), sc2, fma(FLOAT_TYPE(by132.w), sc3, fma(FLOAT_TYPE(by20.w), sc6, FLOAT_TYPE(by232.w) * sc7))))))))))))))); + temp[j][n] = fma(dall, fma(sx, sc0, fma(sy, sc1, fma(sz, sc4, sw * sc5))), fma(-dmin, smin, temp[j][n])); + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp new file mode 100644 index 000000000..0a68891c3 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp @@ -0,0 +1,162 @@ +#version 450 + +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads are used to process each block + const uint it_size = gl_WorkGroupSize.x/16; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid%16; // 0...16 + const uint ix = tid/16; + + const uint il = itid/4; // 0...3 + const uint ir = itid - 4*il; // 0...7 or 0...3 + + const uint v_im = il / 2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224 + const uint v_in = il % 2; + + const uint l0 = 4*ir + 2*v_in; // 0...15 + const uint q_offset = 32*v_im + l0; + const uint y_offset = 64*v_im + l0; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { + const uint y1_idx = i * QUANT_K + y_offset; + const uint y2_idx = y1_idx + 128; + + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; + f16vec2 d = data_a[ib0 + i].d; + const FLOAT_TYPE dall = FLOAT_TYPE(d.x); + const FLOAT_TYPE dmin = FLOAT_TYPE(d.y); + + uint32_t scale0_u32 = data_a_packed16[ib0 + i].scales[v_im ]; + uint32_t scale4_u32 = data_a_packed16[ib0 + i].scales[v_im + 2]; + uint32_t scale8_u32 = data_a_packed16[ib0 + i].scales[v_im + 4]; + uvec4 scale0 = uvec4(unpack8(scale0_u32)); + uvec4 scale4 = uvec4(unpack8(scale4_u32)); + uvec4 scale8 = uvec4(unpack8(scale8_u32)); + + const uint32_t sc0 = ( scale0.x & 0x3f); + const uint32_t sc1 = ( scale0.y & 0x3f); + const uint32_t sc2 = ( scale4.x & 0x3f); + const uint32_t sc3 = ( scale4.y & 0x3f); + const uint32_t sc4 = (( scale8.x & 0x0f) | ((scale0.x & 0xc0) >> 2)); + const uint32_t sc5 = (( scale8.y & 0x0f) | ((scale0.y & 0xc0) >> 2)); + const uint32_t sc6 = (((scale8.x >> 4) & 0x0f) | ((scale4.x & 0xc0) >> 2)); + const uint32_t sc7 = (((scale8.y >> 4) & 0x0f) | ((scale4.y & 0xc0) >> 2)); + + uint32_t qs0_16_u32 = uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 8]) << 16); + uint32_t qs64_80_u32 = uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 32]) | (uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 40]) << 16); + + uint32_t qs0_16_u32_lo4 = qs0_16_u32 & 0x0F0F0F0F; + uint32_t qs0_16_u32_hi4 = (qs0_16_u32 >> 4) & 0x0F0F0F0F; + uint32_t qs64_80_u32_lo4 = qs64_80_u32 & 0x0F0F0F0F; + uint32_t qs64_80_u32_hi4 = (qs64_80_u32 >> 4) & 0x0F0F0F0F; + + uint32_t qh = pack32(u16vec2(data_a_packed16[ib0 + i].qh[l0 / 2], data_a_packed16[ib0 + i].qh[l0 / 2 + 8])); + + uint32_t qs0_16_lo4_offset16 = ((qh >> (2*v_im)) & 0x01010101) << 4; + uint32_t qs0_16_hi4_offset16 = ((qh >> (2*v_im)) & 0x02020202) << 3; + uint32_t qs64_80_lo4_offset16 = ((qh >> (2*v_im)) & 0x10101010) << 0; + uint32_t qs64_80_hi4_offset16 = ((qh >> (2*v_im)) & 0x20202020) >> 1; + + qs0_16_u32_lo4 += qs0_16_lo4_offset16; + qs0_16_u32_hi4 += qs0_16_hi4_offset16; + qs64_80_u32_lo4 += qs64_80_lo4_offset16; + qs64_80_u32_hi4 += qs64_80_hi4_offset16; + + uvec4 qs0_16_lo4 = uvec4(unpack8(qs0_16_u32_lo4)); + uvec4 qs64_80_lo4 = uvec4(unpack8(qs64_80_u32_lo4)); + uvec4 qs0_16_hi4 = uvec4(unpack8(qs0_16_u32_hi4)); + uvec4 qs64_80_hi4 = uvec4(unpack8(qs64_80_u32_hi4)); + + const uint32_t q4_0 = qs0_16_lo4.x; + const uint32_t q4_1 = qs0_16_lo4.y; + const uint32_t q4_2 = qs0_16_lo4.z; + const uint32_t q4_3 = qs0_16_lo4.w; + const uint32_t q4_4 = qs0_16_hi4.x; + const uint32_t q4_5 = qs0_16_hi4.y; + const uint32_t q4_6 = qs0_16_hi4.z; + const uint32_t q4_7 = qs0_16_hi4.w; + const uint32_t q4_8 = qs64_80_lo4.x; + const uint32_t q4_9 = qs64_80_lo4.y; + const uint32_t q4_10 = qs64_80_lo4.z; + const uint32_t q4_11 = qs64_80_lo4.w; + const uint32_t q4_12 = qs64_80_hi4.x; + const uint32_t q4_13 = qs64_80_hi4.y; + const uint32_t q4_14 = qs64_80_hi4.z; + const uint32_t q4_15 = qs64_80_hi4.w; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + B_TYPE_VEC2 by10 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2]; + B_TYPE_VEC2 by116 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 8]; + B_TYPE_VEC2 by132 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 16]; + B_TYPE_VEC2 by148 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 24]; + B_TYPE_VEC2 by20 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2]; + B_TYPE_VEC2 by216 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 8]; + B_TYPE_VEC2 by232 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 16]; + B_TYPE_VEC2 by248 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 24]; + + const FLOAT_TYPE sx = + fma(FLOAT_TYPE(by10.x), q4_0, + fma(FLOAT_TYPE(by10.y), q4_1, + fma(FLOAT_TYPE(by116.x), q4_2, + FLOAT_TYPE(by116.y) * q4_3))); + const FLOAT_TYPE sy = + fma(FLOAT_TYPE(by132.x), q4_4, + fma(FLOAT_TYPE(by132.y), q4_5, + fma(FLOAT_TYPE(by148.x), q4_6, + FLOAT_TYPE(by148.y) * q4_7))); + const FLOAT_TYPE sz = + fma(FLOAT_TYPE(by20.x), q4_8, + fma(FLOAT_TYPE(by20.y), q4_9, + fma(FLOAT_TYPE(by216.x), q4_10, + FLOAT_TYPE(by216.y) * q4_11))); + const FLOAT_TYPE sw = + fma(FLOAT_TYPE(by232.x), q4_12, + fma(FLOAT_TYPE(by232.y), q4_13, + fma(FLOAT_TYPE(by248.x), q4_14, + FLOAT_TYPE(by248.y) * q4_15))); + const FLOAT_TYPE smin = + fma(FLOAT_TYPE(by10.x) + FLOAT_TYPE(by10.y) + FLOAT_TYPE(by116.x) + FLOAT_TYPE(by116.y), sc2, + fma(FLOAT_TYPE(by132.x) + FLOAT_TYPE(by132.y) + FLOAT_TYPE(by148.x) + FLOAT_TYPE(by148.y), sc3, + fma(FLOAT_TYPE(by20.x) + FLOAT_TYPE(by20.y) + FLOAT_TYPE(by216.x) + FLOAT_TYPE(by216.y), sc6, + (FLOAT_TYPE(by232.x) + FLOAT_TYPE(by232.y) + FLOAT_TYPE(by248.x) + FLOAT_TYPE(by248.y)) * sc7))); + temp[j][n] = fma(dall, fma(sx, sc0, fma(sy, sc1, fma(sz, sc4, sw * sc5))), fma(-dmin, smin, temp[j][n])); + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp new file mode 100644 index 000000000..70e13a56b --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp @@ -0,0 +1,112 @@ +#version 450 + +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads are used to process each block + const uint it_size = gl_WorkGroupSize.x/16; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid%16; // 0...16 + const uint ix = tid/16; + + const uint step = 8; + + const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... + const uint v_in = itid - step*v_im; // 0...15 or 0...7 + + const uint l0 = 4 * v_in; // 0, 4, 8, ..., 28 + const uint is = v_in / 4; + + const uint ql_offset = 64*v_im + l0; + const uint qh_offset = 32*v_im + l0; + const uint s_offset = 8*v_im + is; + const uint y_offset = 128*v_im + l0; + + FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { + const uint y_idx = i * QUANT_K + y_offset; + + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; + const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d); + + FLOAT_TYPE scales[4]; + scales[0] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 0]); + scales[1] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 2]); + scales[2] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 4]); + scales[3] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 6]); + + uint32_t ql0_u32 = uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 1]) << 16); + uint32_t ql32_u32 = uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 16]) | (uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 17]) << 16); + + uint32_t ql0_u32_lo4 = ql0_u32 & 0x0F0F0F0F; + uint32_t ql0_u32_hi4 = (ql0_u32 >> 4) & 0x0F0F0F0F; + uint32_t ql32_u32_lo4 = ql32_u32 & 0x0F0F0F0F; + uint32_t ql32_u32_hi4 = (ql32_u32 >> 4) & 0x0F0F0F0F; + + uint32_t qh_u32 = uint32_t(data_a_packed16[ib0 + i].qh[qh_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].qh[qh_offset / 2 + 1]) << 16); + uint32_t qh0_u32 = (qh_u32 & 0x03030303) << 4; + uint32_t qh2_u32 = (qh_u32 & 0x0C0C0C0C) << 2; + uint32_t qh4_u32 = (qh_u32 & 0x30303030) << 0; + uint32_t qh6_u32 = (qh_u32 & 0xC0C0C0C0) >> 2; + + uint32_t q0_u32 = ql0_u32_lo4 | qh0_u32; + uint32_t q1_u32 = ql32_u32_lo4 | qh2_u32; + uint32_t q2_u32 = ql0_u32_hi4 | qh4_u32; + uint32_t q3_u32 = ql32_u32_hi4 | qh6_u32; + + uvec4 q0 = uvec4(unpack8(q0_u32)); + uvec4 q1 = uvec4(unpack8(q1_u32)); + uvec4 q2 = uvec4(unpack8(q2_u32)); + uvec4 q3 = uvec4(unpack8(q3_u32)); + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + B_TYPE_VEC4 by0 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4]; + B_TYPE_VEC4 by32 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 8]; + B_TYPE_VEC4 by64 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 16]; + B_TYPE_VEC4 by96 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 24]; + + FLOAT_TYPE sum = FLOAT_TYPE(0.0); + [[unroll]] for (int l = 0; l < 4; ++l) { + sum = fma(FLOAT_TYPE(by0[l]) * scales[0], FLOAT_TYPE(int8_t(q0[l]) - 32), + fma(FLOAT_TYPE(by32[l]) * scales[1], FLOAT_TYPE(int8_t(q1[l]) - 32), + fma(FLOAT_TYPE(by64[l]) * scales[2], FLOAT_TYPE(int8_t(q2[l]) - 32), + fma(FLOAT_TYPE(by96[l]) * scales[3], FLOAT_TYPE(int8_t(q3[l]) - 32), sum)))); + } + temp[j][n] += sum * d; + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp new file mode 100644 index 000000000..48122cbef --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -0,0 +1,631 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require + +#ifdef FLOAT16 +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#endif + +#ifdef COOPMAT +#extension GL_KHR_cooperative_matrix : enable +#extension GL_KHR_memory_scope_semantics : enable +#extension GL_KHR_shader_subgroup_basic : enable +#endif + +#ifdef MUL_MAT_ID +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#endif + +#include "types.comp" + +#ifndef LOAD_VEC_A +#define LOAD_VEC_A 1 +#endif +#ifndef LOAD_VEC_B +#define LOAD_VEC_B 1 +#endif + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; + +#ifdef MUL_MAT_ID +layout (binding = 3) readonly buffer IDS {int data_ids[];}; +#endif + +layout (push_constant) uniform parameter +{ + uint M; + uint N; + uint K; + uint stride_a; + uint stride_b; + uint stride_d; + + uint batch_stride_a; + uint batch_stride_b; + uint batch_stride_d; + +#ifdef MUL_MAT_ID + uint nei0; + uint nei1; + uint nbi1; + uint ne11; +#else + uint k_split; + uint ne02; + uint ne12; + uint broadcast2; + uint broadcast3; +#endif +} p; + +layout (constant_id = 0) const uint BLOCK_SIZE = 64; +layout (constant_id = 1) const uint BM = 64; +layout (constant_id = 2) const uint BN = 64; +layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working with a quant +layout (constant_id = 4) const uint WM = 32; +layout (constant_id = 5) const uint WN = 32; +layout (constant_id = 6) const uint WMITER = 2; +layout (constant_id = 7) const uint TM = 4; +layout (constant_id = 8) const uint TN = 2; +layout (constant_id = 9) const uint TK = 1; // Only needed for coopmat +layout (constant_id = 10) const uint WARP = 32; + +#ifdef COOPMAT +#define SHMEM_STRIDE (BK + 8) +#else +#define SHMEM_STRIDE (BK + 1) +#endif + +shared FLOAT_TYPE buf_a[BM * SHMEM_STRIDE]; +shared FLOAT_TYPE buf_b[BN * SHMEM_STRIDE]; + +#ifdef MUL_MAT_ID +shared u16vec2 row_ids[3072]; +#endif // MUL_MAT_ID + +#define NUM_WARPS (BLOCK_SIZE / WARP) + +#ifdef COOPMAT +shared ACC_TYPE coopmat_stage[TM * TN * NUM_WARPS]; +#endif + +void main() { +#if defined(DATA_A_IQ4_NL) + init_iq4nl_shmem(); +#endif + +#ifdef MUL_MAT_ID + const uint expert_idx = gl_GlobalInvocationID.z; +#else + const uint batch_idx = gl_GlobalInvocationID.z; + + const uint i13 = batch_idx / p.ne12; + const uint i12 = batch_idx % p.ne12; + + const uint i03 = i13 / p.broadcast3; + const uint i02 = i12 / p.broadcast2; + + const uint batch_idx_a = i03 * p.ne02 + i02; +#endif + + const uint blocks_m = (p.M + BM - 1) / BM; + const uint ir = gl_WorkGroupID.x % blocks_m; + const uint ik = gl_WorkGroupID.x / blocks_m; + const uint ic = gl_WorkGroupID.y; + + const uint WNITER = (WM * WN) / (WARP * TM * TN * WMITER); + const uint WSUBM = WM / WMITER; + const uint WSUBN = WN / WNITER; + +#ifdef COOPMAT + const uint warp_i = gl_SubgroupID; + + const uint tiw = gl_SubgroupInvocationID; + + const uint cms_per_row = WM / TM; + const uint cms_per_col = WN / TN; + + const uint storestride = WARP / TM; + const uint store_r = tiw % TM; + const uint store_c = tiw / TM; +#else + const uint warp_i = gl_LocalInvocationID.x / WARP; + + const uint tiw = gl_LocalInvocationID.x % WARP; + + const uint tiwr = tiw % (WSUBM / TM); + const uint tiwc = tiw / (WSUBM / TM); +#endif + + const uint warp_r = warp_i % (BM / WM); + const uint warp_c = warp_i / (BM / WM); + + const uint loadr_a = gl_LocalInvocationID.x % (BK / LOAD_VEC_A); + const uint loadc_a = gl_LocalInvocationID.x / (BK / LOAD_VEC_A); + const uint loadr_b = gl_LocalInvocationID.x % (BK / LOAD_VEC_B); + const uint loadc_b = gl_LocalInvocationID.x / (BK / LOAD_VEC_B); + + const uint loadstride_a = gl_WorkGroupSize.x * LOAD_VEC_A / BK; + const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B / BK; + +#ifdef MUL_MAT_ID + uint _ne1 = 0; + for (uint ii1 = 0; ii1 < p.nei1; ii1++) { + for (uint ii0 = 0; ii0 < p.nei0; ii0++) { + if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { + row_ids[_ne1] = u16vec2(ii0, ii1); + _ne1++; + } + } + } + + barrier(); + + // Workgroup has no work + if (ic * BN >= _ne1) return; +#endif + +#ifdef MUL_MAT_ID + const uint start_k = 0; + const uint end_k = p.K; +#else + const uint start_k = ik * p.k_split; + const uint end_k = min(p.K, (ik + 1) * p.k_split); +#endif + + uint pos_a = ( +#ifdef MUL_MAT_ID + expert_idx * p.batch_stride_a + +#else + batch_idx_a * p.batch_stride_a + +#endif + ir * BM * p.stride_a + start_k) / LOAD_VEC_A; +#ifdef MUL_MAT_ID + uint pos_b = 0; +#else + uint pos_b = (batch_idx * p.batch_stride_b + ic * BN * p.stride_b + start_k) / LOAD_VEC_B; +#endif + +#ifdef COOPMAT + coopmat cache_a; + coopmat cache_b; + coopmat sums[cms_per_row * cms_per_col]; + + [[unroll]] for (uint i = 0; i < cms_per_row * cms_per_col; i++) { + sums[i] = coopmat(0.0f); + } +#else + ACC_TYPE sums[WMITER * TM * WNITER * TN]; + FLOAT_TYPE cache_a[WMITER * TM]; + FLOAT_TYPE cache_b[WNITER * TN]; + + [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN; i++) { + sums[i] = ACC_TYPE(0.0f); + } +#endif + + for (uint block = start_k; block < end_k; block += BK) { + [[unroll]] for (uint l = 0; l < BM; l += loadstride_a) { + +#if defined(DATA_A_F32) || defined(DATA_A_F16) +#if LOAD_VEC_A == 8 + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx][0].x); + buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx][0].y); + buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx][0].z); + buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx][0].w); + buf_a[buf_idx + 4] = FLOAT_TYPE(data_a[idx][1].x); + buf_a[buf_idx + 5] = FLOAT_TYPE(data_a[idx][1].y); + buf_a[buf_idx + 6] = FLOAT_TYPE(data_a[idx][1].z); + buf_a[buf_idx + 7] = FLOAT_TYPE(data_a[idx][1].w); +#elif LOAD_VEC_A == 4 + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx].x); + buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx].y); + buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx].z); + buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx].w); +#else + if (ir * BM + loadc_a + l < p.M && block + loadr_a < end_k) { + buf_a[(loadc_a + l) * SHMEM_STRIDE + loadr_a] = FLOAT_TYPE(data_a[pos_a + (loadc_a + l) * p.stride_a + loadr_a]); + } else { + buf_a[(loadc_a + l) * SHMEM_STRIDE + loadr_a] = FLOAT_TYPE(0.0f); + } +#endif +#elif defined(DATA_A_Q4_0) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; + + const uint ib = idx / 16; + const uint iqs = idx & 0xF; + + const float d = float(data_a[ib].d); + const uint vui = uint(data_a[ib].qs[iqs]); + const vec2 v = (vec2(vui & 0xF, vui >> 4) - 8.0f) * d; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q4_1) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; + + const uint ib = idx / 16; + const uint iqs = idx & 0xF; + + const float d = float(data_a[ib].d); + const float m = float(data_a[ib].m); + const uint vui = uint(data_a[ib].qs[iqs]); + const vec2 v = vec2(vui & 0xF, vui >> 4) * d + m; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q5_0) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; + + const uint ib = idx / 16; + const uint iqs = idx & 0xF; + + const float d = float(data_a[ib].d); + const uint uint_qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0]; + const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const uint vui = uint(data_a[ib].qs[iqs]); + const vec2 v = (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f) * d; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q5_1) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; + + const uint ib = idx / 16; + const uint iqs = idx & 0xF; + + const float d = float(data_a[ib].d); + const float m = float(data_a[ib].m); + const uint uint_qh = data_a[ib].qh; + const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); + const uint vui = uint(data_a[ib].qs[iqs]); + const vec2 v = vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) * d + m; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q8_0) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 16; + const uint iqs = (idx & 0xF) * 2; + + const float d = float(data_a[ib].d); + const vec2 v = vec2(int(data_a[ib].qs[iqs]), int(data_a[ib].qs[iqs + 1])) * d; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 1] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q2_K) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint qsi = (iqs / 64) * 32 + (iqs % 16) * 2; // 0,2,4..30 + const uint scalesi = iqs / 8; // 0..15 + const uint qsshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 + + const uvec2 qs = uvec2(data_a[ib].qs[qsi], data_a[ib].qs[qsi + 1]); + const uint scales = data_a[ib].scales[scalesi]; + const vec2 d = vec2(data_a[ib].d); + + const vec2 v = d.x * float(scales & 0xF) * vec2((qs >> qsshift) & 3) - d.y * float(scales >> 4); + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 1] = FLOAT_TYPE(v.y); +#elif defined(DATA_A_Q3_K) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint n = iqs / 64; // 0,1 + const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 + const uint hmi = (iqs % 16) * 2; // 0,2,4..30 + const uint j = (iqs % 64) / 4; // 0..3 + const uint is = iqs / 8; // 0..15 + const uint halfsplit = ((iqs % 64) / 16); // 0,1,2,3 + const uint qsshift = halfsplit * 2; // 0,2,4,6 + const uint m = 1 << (4 * n + halfsplit); // 1,2,4,8,16,32,64,128 + + const int8_t us = int8_t(is < 4 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+8] >> 0) & 3) << 4) : + is < 8 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+4] >> 2) & 3) << 4) : + is < 12 ? (data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is+0] >> 4) & 3) << 4) : + (data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is-4] >> 6) & 3) << 4)); + const float dl = float(data_a[ib].d) * float(us - 32); + + buf_a[buf_idx ] = FLOAT_TYPE(dl * float(int8_t((data_a[ib].qs[qsi ] >> qsshift) & 3) - (((data_a[ib].hmask[hmi ] & m) != 0) ? 0 : 4))); + buf_a[buf_idx + 1] = FLOAT_TYPE(dl * float(int8_t((data_a[ib].qs[qsi + 1] >> qsshift) & 3) - (((data_a[ib].hmask[hmi + 1] & m) != 0) ? 0 : 4))); +#elif defined(DATA_A_Q4_K) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint n = iqs / 32; // 0,1,2,3 + const uint b = (iqs % 32) / 16; // 0,1 + const uint is = 2 * n + b; // 0..7 + const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 + + const vec2 loadd = vec2(data_a[ib].d); + + const uint scidx0 = (is < 4) ? is : (is + 4); + const uint scidx1 = (is < 4) ? is : (is - 4); + const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; + const uint scidxshift1 = (is < 4) ? 0 : 2; + const uint mbidx0 = is + 4; + const uint mbidx1 = (is < 4) ? is + 4 : is; + const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; + const uint mbidxshift0 = (is < 4) ? 0 : 4; + const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + const uint mbidxshift1 = (is < 4) ? 0 : 2; + + const uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + const uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const float d = loadd.x * sc; + const float m = -loadd.y * mbyte; + + buf_a[buf_idx ] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi ] >> (b * 4)) & 0xF), m)); + buf_a[buf_idx + 1] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi + 1] >> (b * 4)) & 0xF), m)); +#elif defined(DATA_A_Q5_K) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint n = iqs / 32; // 0,1,2,3 + const uint b = (iqs % 32) / 16; // 0,1 + const uint is = 2 * n + b; // 0..7 + const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 + const uint qhi = (iqs % 16) * 2; // 0,2,4..30 + + const uint8_t hm = uint8_t(1 << (iqs / 16)); + + const vec2 loadd = vec2(data_a[ib].d); + + const uint scidx0 = (is < 4) ? is : (is + 4); + const uint scidx1 = (is < 4) ? is : (is - 4); + const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; + const uint scidxshift1 = (is < 4) ? 0 : 2; + const uint mbidx0 = is + 4; + const uint mbidx1 = (is < 4) ? is + 4 : is; + const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; + const uint mbidxshift0 = (is < 4) ? 0 : 4; + const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; + const uint mbidxshift1 = (is < 4) ? 0 : 2; + + const uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); + const uint8_t mbyte = uint8_t(((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); + + const float d = loadd.x * sc; + const float m = -loadd.y * mbyte; + + buf_a[buf_idx ] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi ] >> (b * 4)) & 0xF) + float((data_a[ib].qh[qhi ] & hm) != 0 ? 16 : 0), m)); + buf_a[buf_idx + 1] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi + 1] >> (b * 4)) & 0xF) + float((data_a[ib].qh[qhi + 1] & hm) != 0 ? 16 : 0), m)); +#elif defined(DATA_A_Q6_K) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint n = iqs / 64; // 0,1 + const uint b = (iqs % 64) / 32; // 0,1 + const uint is_b = (iqs % 16) / 8; // 0,1 + const uint qhshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 + const uint is = 8 * n + qhshift + is_b; // 0..15 + const uint qsi = n * 64 + (iqs % 32) * 2; // 0,2,4..126 + const uint qhi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 + + const float dscale = float(data_a[ib].d) * float(data_a[ib].scales[is]); + + buf_a[buf_idx ] = FLOAT_TYPE(dscale * float(int8_t(((data_a[ib].ql[qsi ] >> (b * 4)) & 0xF) | (((data_a[ib].qh[qhi ] >> qhshift) & 3) << 4)) - 32)); + buf_a[buf_idx + 1] = FLOAT_TYPE(dscale * float(int8_t(((data_a[ib].ql[qsi + 1] >> (b * 4)) & 0xF) | (((data_a[ib].qh[qhi + 1] >> qhshift) & 3) << 4)) - 32)); +#elif defined(DATA_A_IQ4_NL) + const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; + const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; + + const uint ib = idx / 16; + const uint iqs = idx & 0xF; + + const float d = float(data_a[ib].d); + const uint vui = uint(data_a[ib].qs[iqs]); + const vec2 v = vec2(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[vui >> 4]) * d; + + buf_a[buf_idx ] = FLOAT_TYPE(v.x); + buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); +#endif + } + [[unroll]] for (uint l = 0; l < BN; l += loadstride_b) { +#if LOAD_VEC_B == 8 +#ifdef MUL_MAT_ID + const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l]; + const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + loadr_b; +#else + const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b; +#endif + const uint buf_idx = (loadc_b + l) * SHMEM_STRIDE + loadr_b * LOAD_VEC_B; + buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx][0].x); + buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx][0].y); + buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx][0].z); + buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx][0].w); + buf_b[buf_idx + 4] = FLOAT_TYPE(data_b[idx][1].x); + buf_b[buf_idx + 5] = FLOAT_TYPE(data_b[idx][1].y); + buf_b[buf_idx + 6] = FLOAT_TYPE(data_b[idx][1].z); + buf_b[buf_idx + 7] = FLOAT_TYPE(data_b[idx][1].w); +#elif LOAD_VEC_B == 4 +#ifdef MUL_MAT_ID + const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l]; + const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + loadr_b; +#else + const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b; +#endif + const uint buf_idx = (loadc_b + l) * SHMEM_STRIDE + loadr_b * LOAD_VEC_B; + buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx].x); + buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx].y); + buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx].z); + buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx].w); +#elif !MUL_MAT_ID + if (ic * BN + loadc_b + l < p.N && block + loadr_b < end_k) { + buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(data_b[pos_b + (loadc_b + l) * p.stride_b + loadr_b]); + } else { + buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(0.0f); + } +#else + const uint row_i = ic * BN + loadc_b + l; + if (row_i < _ne1) { + const u16vec2 row_idx = row_ids[row_i]; + buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(data_b[pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + loadr_b]); + } else { + buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(0.0f); + } +#endif + } + + barrier(); + + pos_a += BK / LOAD_VEC_A; + pos_b += BK / LOAD_VEC_B; + +#ifdef COOPMAT + [[unroll]] for (uint i = 0; i < BK; i += TK) { + [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { + // Load from shared into cache + coopMatLoad(cache_a, buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { + coopMatLoad(cache_b, buf_b, (warp_c * WN + cm_col * TN) * SHMEM_STRIDE + i, SHMEM_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + + sums[cm_col * cms_per_row + cm_row] = coopMatMulAdd(cache_a, cache_b, sums[cm_col * cms_per_row + cm_row]); + } + } + } +#else + [[unroll]] for (uint i = 0; i < BK; i++) { + // Load from shared into cache + [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { + [[unroll]] for (uint j = 0; j < TM; j++) { + cache_a[wsir * TM + j] = buf_a[(warp_r * WM + wsir * WSUBM + tiwr * TM + j) * SHMEM_STRIDE + i]; + } + } + [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { + [[unroll]] for (uint j = 0; j < TN; j++) { + cache_b[wsic * TN + j] = buf_b[(warp_c * WN + wsic * WSUBN + tiwc * TN + j) * SHMEM_STRIDE + i]; + } + } + + [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { + [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { + [[unroll]] for (uint cc = 0; cc < TN; cc++) { + [[unroll]] for (uint cr = 0; cr < TM; cr++) { + const uint sums_idx = (wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr; + sums[sums_idx] = fma(ACC_TYPE(cache_a[wsir * TM + cr]), ACC_TYPE(cache_b[wsic * TN + cc]), sums[sums_idx]); + } + } + } + } + } +#endif + + barrier(); + } + + const uint dr = ir * BM + warp_r * WM; + const uint dc = ic * BN + warp_c * WN; + +#ifndef MUL_MAT_ID + const uint offsets = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * gl_NumWorkGroups.z; +#endif + +#ifdef COOPMAT +#ifdef MUL_MAT_ID + [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { + [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { + coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); + + [[unroll]] for (uint col = 0; col < BN; col += storestride) { + const uint row_i = dc + cm_col * TN + col + store_c; + if (row_i >= _ne1) break; + + const u16vec2 row_idx = row_ids[row_i]; + + data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); + } + } + } +#else + const bool is_aligned = p.stride_d % 4 == 0; // Assumption: D_TYPE == float + + [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { + [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { + const bool is_in_bounds = dr + (cm_row + 1) * TM <= p.M && dc + (cm_col + 1) * TN <= p.N; + + if (is_aligned && is_in_bounds) { + // Full coopMat is within bounds and stride_d is aligned with 16B + coopmat cm_dtype = coopmat(sums[cm_col * cms_per_row + cm_row]); + coopMatStore(cm_dtype, data_d, offsets + (dc + cm_col * TN) * p.stride_d + dr + cm_row * TM, p.stride_d, gl_CooperativeMatrixLayoutColumnMajor); + } else if (is_in_bounds) { + // Full coopMat is within bounds, but stride_d is not aligned + coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); + + [[unroll]] for (uint col = 0; col < TN; col += storestride) { + data_d[offsets + (dc + cm_col * TN + col + store_c) * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); + } + } else if (dr + cm_row * TM < p.M && dc + cm_col * TN < p.N) { + // Partial coopMat is within bounds + coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); + + [[unroll]] for (uint col = 0; col < TN; col += storestride) { + if (dr + cm_row * TM + store_r < p.M && dc + cm_col * TN + col + store_c < p.N) { + data_d[offsets + (dc + cm_col * TN + col + store_c) * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); + } + } + } + } + } +#endif // MUL_MAT_ID +#else + [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { + [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { + + const uint dr_warp = dr + wsir * WSUBM + tiwr * TM; + const uint dc_warp = dc + wsic * WSUBN + tiwc * TN; + [[unroll]] for (uint cc = 0; cc < TN; cc++) { +#ifdef MUL_MAT_ID + const uint row_i = dc_warp + cc; + if (row_i >= _ne1) break; + + const u16vec2 row_idx = row_ids[row_i]; +#endif // MUL_MAT_ID + [[unroll]] for (uint cr = 0; cr < TM; cr++) { +#ifdef MUL_MAT_ID + data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]); +#else + if (dr_warp + cr < p.M && dc_warp + cc < p.N) { + data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]); + } +#endif // MUL_MAT_ID + } + } + } + } +#endif // COOPMAT +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp new file mode 100644 index 000000000..cbfa5dce1 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp @@ -0,0 +1,328 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require + +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +#extension GL_KHR_memory_scope_semantics : enable +#extension GL_KHR_cooperative_matrix : enable +#extension GL_NV_cooperative_matrix2 : enable +#extension GL_EXT_buffer_reference : enable +#extension GL_KHR_shader_subgroup_ballot : enable +#extension GL_KHR_shader_subgroup_vote : enable + +#include "types.comp" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (constant_id = 1) const uint BM = 64; +layout (constant_id = 2) const uint BN = 64; +layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working with a quant + +layout (push_constant) uniform parameter +{ + uint M; + uint N; + uint K; + uint stride_a; + uint stride_b; + uint stride_d; + + uint batch_stride_a; + uint batch_stride_b; + uint batch_stride_d; + +#ifdef MUL_MAT_ID + uint nei0; + uint nei1; + uint nbi1; + uint ne11; +#else + uint k_split; + uint ne02; + uint ne12; + uint broadcast2; + uint broadcast3; +#endif +} p; + + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; + +#if QUANT_K > 1 +#define DECODEFUNCA , dequantFuncA +#define MAT_A_TYPE float16_t + +#include "dequant_funcs_cm2.comp" + +#else +#define DECODEFUNCA +#define MAT_A_TYPE A_TYPE +#endif + +#define MAT_B_TYPE B_TYPE + +#ifdef MUL_MAT_ID +layout (binding = 3) readonly buffer IDS {int data_ids[];}; + +shared u16vec4 row_ids[3072]; + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufB { + B_TYPE b[]; +}; + +uint _ne1; +shared uint _ne1_sh; + +B_TYPE decodeFuncB(const in decodeBufB bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint row_i = blockCoords[0]; + + if (row_i >= _ne1) { + return B_TYPE(0.0); + } + + const u16vec4 row_idx = row_ids[row_i]; + B_TYPE ret = data_b[row_idx.y * p.batch_stride_b + row_idx.x * p.stride_b + blockCoords[1]]; + + return ret; +} + +D_TYPE perElemOpD(const in uint32_t r, const in uint32_t c, const in D_TYPE elem, const in uint32_t ir, const in uint32_t ic) +{ + uint dr = ir * BM + r; + uint dc = ic * BN + c; + + if (dr < p.M && dc < _ne1) { + uint row_i = dc; + const u16vec4 row_idx = row_ids[row_i]; + data_d[row_idx.y * p.batch_stride_d + row_idx.z * p.stride_d + dr] = elem; + } + return elem; +} + +#endif + +void main() { +#if defined(DATA_A_IQ4_NL) + init_iq4nl_shmem(); +#endif + +#ifdef MUL_MAT_ID + const uint expert_idx = gl_GlobalInvocationID.z; +#else + const uint batch_idx = gl_GlobalInvocationID.z; + + const uint i13 = batch_idx / p.ne12; + const uint i12 = batch_idx % p.ne12; + + const uint i03 = i13 / p.broadcast3; + const uint i02 = i12 / p.broadcast2; + + const uint batch_idx_a = i03 * p.ne02 + i02; +#endif + + const uint blocks_m = (p.M + BM - 1) / BM; + const uint ir = gl_WorkGroupID.x % blocks_m; + const uint ik = gl_WorkGroupID.x / blocks_m; + const uint ic = gl_WorkGroupID.y; + +#ifdef MUL_MAT_ID + // Spread the search across all elements in the first subgroup + if (gl_SubgroupID == 0) { + _ne1 = 0; + uint num_elements = p.nei1 * p.nei0; + + for (uint i = gl_SubgroupInvocationID; subgroupAny(i < num_elements); i += gl_SubgroupSize) { + bool in_range = i < num_elements; + uint ii0 = i % p.nei0; + uint ii1 = i / p.nei0; + uint id = in_range ? data_ids[ii1*p.nbi1 + ii0] : 0; + uvec4 ballot = subgroupBallot(in_range && id == expert_idx); + uint idx = subgroupBallotExclusiveBitCount(ballot); + if (in_range && id == expert_idx) { + row_ids[_ne1 + idx] = u16vec4(ii0 % p.ne11, ii1, ii0, 0); + } + _ne1 += subgroupBallotBitCount(ballot); + } + _ne1_sh = _ne1; + } + + barrier(); + + _ne1 = _ne1_sh; + + // Workgroup has no work + if (ic * BN >= _ne1) return; +#endif + +#ifdef MUL_MAT_ID + uint start_k = 0; + const uint end_k = p.K; +#else + uint start_k = ik * p.k_split; + const uint end_k = min(p.K, (ik + 1) * p.k_split); +#endif + + coopmat sum; + sum = coopmat(0.0); + +#ifdef MUL_MAT_ID + uint pos_a = (expert_idx * p.batch_stride_a) / QUANT_K; + uint pos_b = 0; +#else + uint pos_a = (batch_idx_a * p.batch_stride_a) / QUANT_K; + uint pos_b = batch_idx * p.batch_stride_b; +#endif + + uint stride_a = p.stride_a / QUANT_K; + uint stride_b = p.stride_b; + + // Hint to the compiler that values are aligned (want 16B alignment). + // Quants are always block-aligned, no alignment needed. +#if ALIGNED +#if QUANT_K == 1 + stride_a &= ~7; +#endif + stride_b &= ~7; +#endif + + // Create layouts for both clamped and unclamped accesses + tensorLayoutNV<2> tensorLayoutA = createTensorLayoutNV(2); + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutAClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutNV<2> tensorLayoutB = createTensorLayoutNV(2); + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutBClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + +#if QUANT_K > 1 + tensorLayoutA = setTensorLayoutBlockSizeNV(tensorLayoutA, 1, QUANT_K); + tensorLayoutAClamp = setTensorLayoutBlockSizeNV(tensorLayoutAClamp, 1, QUANT_K); +#endif + + // Use end_k rather than p.K as the dimension because that's what + // we need to bound check against when using split_k + tensorLayoutA = setTensorLayoutDimensionNV(tensorLayoutA, p.M, end_k); + tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, p.N, end_k); + tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.N, p.M); + tensorLayoutAClamp = setTensorLayoutDimensionNV(tensorLayoutAClamp, p.M, end_k); + tensorLayoutBClamp = setTensorLayoutDimensionNV(tensorLayoutBClamp, p.N, end_k); + + tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0); + +#if !defined(MUL_MAT_ID) + // Detect a fast path where all loads are entirely in bounds and no clamping is required + if ((ir + 1) * BM <= p.M && (ic + 1) * BN <= p.N && (start_k % BK) == 0 && (end_k % BK) == 0 && +#if QUANT_K == 1 + (stride_a % 8) == 0 && +#endif + (stride_b % 8) == 0 && (start_k % 8) == 0) { + // Hint to the compiler that values are aligned (want 16B alignment) + start_k &= ~7; + stride_b &= ~7; +#if QUANT_K == 1 + stride_a &= ~7; +#endif + + tensorLayoutA = setTensorLayoutStrideNV(tensorLayoutA, stride_a, 1); + tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1); + + uint k_iters = (end_k - start_k + BK - 1) / BK; + + for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) { + + coopmat mat_a; + coopmat mat_b; + + coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA); + coopmat mat_a_ft = coopmat(mat_a); + + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose); + coopmat mat_b_ft = coopmat(mat_b); + + sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); + } + } else +#endif // !defined(MUL_MAT_ID) + { + tensorLayoutA = setTensorLayoutStrideNV(tensorLayoutA, stride_a, 1); + + tensorLayoutAClamp = setTensorLayoutStrideNV(tensorLayoutAClamp, stride_a, 1); + + tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1); + + tensorLayoutBClamp = setTensorLayoutStrideNV(tensorLayoutBClamp, stride_b, 1); + + [[dont_unroll]] + for (uint block_k = start_k; block_k < end_k; block_k += BK) { + + coopmat mat_a; + coopmat mat_b; + coopmat mat_a_ft; + coopmat mat_b_ft; + + // Clamping is expensive, so detect different code paths for each combination + // of A and B needing clamping. + bool unclampedA = (ir + 1) * BM <= p.M && block_k + BK <= end_k && (block_k % 8) == 0; +#ifdef MUL_MAT_ID + bool unclampedB = true; +#else + bool unclampedB = (ic + 1) * BN <= p.N && block_k + BK <= end_k && (block_k % 8) == 0; +#endif + if (unclampedA && unclampedB) { + coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, (block_k & ~7), BK) DECODEFUNCA); +#ifdef MUL_MAT_ID + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose, decodeFuncB); +#else + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, (block_k & ~7), BK), tensorViewTranspose); +#endif + mat_a_ft = coopmat(mat_a); + mat_b_ft = coopmat(mat_b); + sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); + } else if (unclampedA && !unclampedB) { + coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, (block_k & ~7), BK) DECODEFUNCA); + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutBClamp, ic * BN, BN, block_k, BK), tensorViewTranspose); + + mat_a_ft = coopmat(mat_a); + mat_b_ft = coopmat(mat_b); + sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); + } else if (!unclampedA && unclampedB) { + coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutAClamp, ir * BM, BM, block_k, BK) DECODEFUNCA); +#ifdef MUL_MAT_ID + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose, decodeFuncB); +#else + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, (block_k & ~7), BK), tensorViewTranspose); +#endif + mat_a_ft = coopmat(mat_a); + mat_b_ft = coopmat(mat_b); + sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); + } else if (!unclampedA && !unclampedB) { + coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutAClamp, ir * BM, BM, block_k, BK) DECODEFUNCA); + coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutBClamp, ic * BN, BN, block_k, BK), tensorViewTranspose); + + mat_a_ft = coopmat(mat_a); + mat_b_ft = coopmat(mat_b); + sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); + } + } + } + + // Convert from ACC_TYPE to D_TYPE + coopmat mat_d; + mat_d = coopmat(sum); + +#ifdef MUL_MAT_ID + // Call callback to store each element, remapping row through shared memory + coopMatPerElementNV(mat_d, mat_d, perElemOpD, ir, ic); +#else + tensorLayoutD = setTensorLayoutStrideNV(tensorLayoutD, p.stride_d, 1); + + uint pos_d = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * gl_NumWorkGroups.z; + coopMatStoreTensorNV(mat_d, data_d, pos_d, sliceTensorLayoutNV(tensorLayoutD, ic * BN, BN, ir * BM, BM), tensorViewTranspose); +#endif +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp new file mode 100644 index 000000000..6627a50bd --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp @@ -0,0 +1,44 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable +#define BLOCK_SIZE 512 + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +shared vec2 sum[BLOCK_SIZE]; + +void main() { + const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationID.x; + + sum[tid] = vec2(0.0f, 0.0f); + + [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { + const float xi = float(data_a[row*p.KX + col]); + sum[tid].x += xi; + sum[tid].y += xi * xi; + } + + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + sum[tid] += sum[tid + s]; + } + barrier(); + } + + const float mean = sum[0].x / p.KX; + const float var = sum[0].y / p.KX - mean * mean; + const float inv_std = inversesqrt(var + p.param1); + + [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { + data_d[row*p.KX + col] = D_TYPE((float(data_a[row*p.KX + col]) - mean) * inv_std); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp new file mode 100644 index 000000000..450b67fc5 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp @@ -0,0 +1,28 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (idx >= p.ne) { + return; + } + + const uint i3 = idx / (p.ne12*p.ne11*p.ne10); + const uint i3_offset = i3 * p.ne12*p.ne11*p.ne10; + const uint i2 = (idx - i3_offset) / (p.ne11*p.ne10); + const uint i2_offset = i2*p.ne11*p.ne10; + const uint i1 = (idx - i3_offset - i2_offset) / p.ne10; + const uint i0 = idx - i3_offset - i2_offset - i1*p.ne10; + + const uint src0_idx = i3*p.nb03 + i2*p.nb02 + i1*p.nb01 + i0*p.nb00; + const uint dst_idx = i3*p.nb13 + i2*p.nb12 + i1*p.nb11 + i0*p.nb10; + + const bool is_src0 = i0 < p.ne00 && i1 < p.ne01 && i2 < p.ne02 && i3 < p.ne03; + + data_d[get_doffset() + dst_idx] = D_TYPE(is_src0 ? data_a[get_aoffset() + src0_idx] : 0.0f); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp new file mode 100644 index 000000000..b6124411a --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp @@ -0,0 +1,74 @@ +#version 450 + +#include "types.comp" + +#extension GL_EXT_shader_16bit_storage : require + +layout(push_constant) uniform parameter { + uint IW; uint IH; + uint OW; uint OH; + uint OC; + uint pelements; + uint op; + int k0; int k1; + int s0; int s1; + int p0; int p1; +} p; + +#define BLOCK_SIZE 512 +#define FLT_MAX 3.402823466e+38F +#define OP_POOL_MAX 0u +#define OP_POOL_AVG 1u + +layout (local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout(binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint idx = gl_GlobalInvocationID.x; + if (idx >= p.pelements) { + return; + } + + const uint O_HW = p.OW * p.OH; + + const uint nc = idx / O_HW; + const uint cur_oh = (idx % O_HW) / p.OW; + const uint cur_ow = (idx % O_HW) % p.OW; + + const int start_h = int(cur_oh) * p.s0 - p.p0; + const uint bh = max(start_h, 0); + const uint eh = min(start_h + p.k0, p.IH); + + const int start_w = int(cur_ow) * p.s1 - p.p1; + const uint bw = max(start_w, 0); + const uint ew = min(start_w + p.k1, p.IW); + + const float scale = 1.0 / float(p.k0 * p.k1); + float res; + + if (p.op == OP_POOL_AVG) { + res = 0.0; + } else if (p.op == OP_POOL_MAX) { + res = -FLT_MAX; + } else { + return; + } + + #pragma unroll + for (uint i = bh; i < eh; i++) { + #pragma unroll + for (uint j = bw; j < ew; j++) { + const float cur = D_TYPE(data_a[nc * p.IH * p.IW + i * p.IW + j]); + + if (p.op == OP_POOL_AVG) { + res += cur * scale; + } else if (p.op == OP_POOL_MAX) { + res = max(res, cur); + } + } + } + + data_d[nc * O_HW + cur_oh * p.OW + cur_ow] = res; +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp new file mode 100644 index 000000000..52a19b62a --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp @@ -0,0 +1,21 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + + data_d[i] = max(float(data_a[i]), 0); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp new file mode 100644 index 000000000..1568b141d --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp @@ -0,0 +1,26 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +uint src0_idx_mod(uint idx) { + const uint i13 = idx / (p.ne12*p.ne11*p.ne10); + const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10; + const uint i12 = (idx - i13_offset) / (p.ne11*p.ne10); + const uint i12_offset = i12*p.ne11*p.ne10; + const uint i11 = (idx - i13_offset - i12_offset) / p.ne10; + const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10; + return (i13 % p.ne03)*p.nb03 + (i12 % p.ne02)*p.nb02 + (i11 % p.ne01)*p.nb01 + (i10 % p.ne00)*p.nb00; +} + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(data_a[get_aoffset() + src0_idx_mod(idx)]); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp new file mode 100644 index 000000000..b554400ba --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp @@ -0,0 +1,42 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable +#define BLOCK_SIZE 512 + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +shared FLOAT_TYPE sum[BLOCK_SIZE]; + +void main() { + const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationID.x; + + sum[tid] = FLOAT_TYPE(0.0f); // partial sum for thread in warp + + [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { + const FLOAT_TYPE xi = FLOAT_TYPE(data_a[row*p.KX + col]); + sum[tid] += xi * xi; + } + + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + sum[tid] += sum[tid + s]; + } + barrier(); + } + + const FLOAT_TYPE mean = sum[0] / FLOAT_TYPE(p.KX); + const FLOAT_TYPE scale = inversesqrt(mean + FLOAT_TYPE(p.param1)); + + [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { + data_d[row*p.KX + col] = D_TYPE(scale * FLOAT_TYPE(data_a[row*p.KX + col])); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp new file mode 100644 index 000000000..574b51ca5 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp @@ -0,0 +1,49 @@ +#include "types.comp" + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_spirv_intrinsics: enable + +#if RTE16 +spirv_execution_mode(capabilities = [4467], 4462, 16); // RoundingModeRTE, 16 bits +#endif + +layout(local_size_x = 1, local_size_y = 256, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer Y {int data_pos[];}; +layout (binding = 2) readonly buffer Z {float data_ff[];}; +layout (binding = 3) writeonly buffer D {D_TYPE data_d[];}; + +layout (push_constant) uniform parameter { + uint ncols; + uint n_dims; + float freq_scale; + uint p_delta_rows; + float freq_base; + float ext_factor; + float attn_factor; + float corr_dims[2]; + float theta_scale; + uint has_ff; +} p; + +float rope_yarn_ramp(const float low, const float high, const uint i0) { + const float y = (i0 / 2 - low) / max(0.001f, high - low); + return 1.0f - min(1.0f, max(0.0f, y)); +} + +void rope_yarn(const float theta_extrap, const uint i0, out float cos_theta, out float sin_theta) { + float mscale = p.attn_factor; + // Get n-d rotational scaling corrected for extrapolation + float theta_interp = p.freq_scale * theta_extrap; + float theta = theta_interp; + if (p.ext_factor != 0.0f) { + float ramp_mix = rope_yarn_ramp(p.corr_dims[0], p.corr_dims[1], i0) * p.ext_factor; + theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix; + + // Get n-d magnitude scaling corrected for interpolation + mscale *= 1.0f + 0.1f * log(1.0f / p.freq_scale); + } + cos_theta = cos(theta) * mscale; + sin_theta = sin(theta) * mscale; +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp new file mode 100644 index 000000000..83b46b69b --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp @@ -0,0 +1,37 @@ +#version 450 + +#include "rope_head.comp" + +void main() { + const uint col = gl_GlobalInvocationID.y * 2; + const uint row = gl_GlobalInvocationID.x; + + if (col >= p.ncols) { + return; + } + + if (col >= p.n_dims) { + const uint i = row*p.ncols + col; + + data_d[i + 0] = data_a[i + 0]; + data_d[i + 1] = data_a[i + 1]; + + return; + } + + const uint i = row*p.ncols + col/2; + const uint i2 = row/p.p_delta_rows; + + const float theta_base = data_pos[i2] * pow(p.theta_scale, col/2.0f); + + const float freq_factor = p.has_ff != 0 ? data_ff[col/2] : 1.0f; + + float cos_theta, sin_theta; + rope_yarn(theta_base / freq_factor, col, cos_theta, sin_theta); + + const float x0 = float(data_a[i + 0]); + const float x1 = float(data_a[i + p.n_dims/2]); + + data_d[i + 0] = D_TYPE(x0*cos_theta - x1*sin_theta); + data_d[i + p.n_dims/2] = D_TYPE(x0*sin_theta + x1*cos_theta); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp new file mode 100644 index 000000000..e416ad938 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp @@ -0,0 +1,37 @@ +#version 450 + +#include "rope_head.comp" + +void main() { + const uint col = gl_GlobalInvocationID.y * 2; + const uint row = gl_GlobalInvocationID.x; + + if (col >= p.ncols) { + return; + } + + if (col >= p.n_dims) { + const uint i = row*p.ncols + col; + + data_d[i + 0] = data_a[i + 0]; + data_d[i + 1] = data_a[i + 1]; + + return; + } + + const uint i = row*p.ncols + col; + const uint i2 = row/p.p_delta_rows; + + const float theta_base = data_pos[i2] * pow(p.theta_scale, col/2.0f); + + const float freq_factor = p.has_ff != 0 ? data_ff[col/2] : 1.0f; + + float cos_theta, sin_theta; + rope_yarn(theta_base / freq_factor, col, cos_theta, sin_theta); + + const float x0 = float(data_a[i + 0]); + const float x1 = float(data_a[i + 1]); + + data_d[i + 0] = D_TYPE(x0*cos_theta - x1*sin_theta); + data_d[i + 1] = D_TYPE(x0*sin_theta + x1*cos_theta); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp new file mode 100644 index 000000000..4663428de --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp @@ -0,0 +1,24 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +const uint num_threads = 128; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + uint idx = get_idx(); + + // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation + const uint num_iter = 4; + + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (idx >= p.ne) { + continue; + } + + data_d[get_doffset() + idx] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + idx]) * FLOAT_TYPE(p.param1)); + idx += num_threads; + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp new file mode 100644 index 000000000..4d36f88e0 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp @@ -0,0 +1,22 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + + const float xi = float(data_a[i]); + data_d[i] = D_TYPE(xi / (1.0f + exp(-xi))); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp new file mode 100644 index 000000000..d7c15a169 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp @@ -0,0 +1,17 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + + const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(sin(val)); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp new file mode 100644 index 000000000..a25808e16 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp @@ -0,0 +1,174 @@ +#version 450 + +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +layout (push_constant) uniform parameter +{ + uint KX; + uint KY; + float scale; + float max_bias; + float m0; + float m1; + uint n_head_log2; + uint nrows_x; +} p; + +#include "types.comp" + +layout(constant_id = 0) const uint BLOCK_SIZE = 32; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer Y {B_TYPE data_b[];}; +layout (binding = 2) buffer D {D_TYPE data_d[];}; + +shared FLOAT_TYPE vals[BLOCK_SIZE]; + +// num_iters is the number of BLOCK_SIZE loop iterations we need to iterate +// over all the columns. The main function tries to pass a constant here, +// as if it were a template function, to allow unrolling. +void soft_max(uint num_iters) { + const uint tid = gl_LocalInvocationID.x; + const uint rowx = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; + const uint rowy = (p.KY > 0) ? (rowx % p.KY) : 0; + + if (rowx >= p.nrows_x) { + return; + } + + float slope = 1.0f; + + // ALiBi + if (p.max_bias > 0.0f) { + const uint h = rowx/p.KY; // head index + + const float base = h < p.n_head_log2 ? p.m0 : p.m1; + const uint exp = h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1; + + slope = pow(base, exp); + } + + // Find max + FLOAT_TYPE max_val = uintBitsToFloat(0xFF800000); + + // Cache values while we compute the max, so we don't need to read them + // again when we're ready to compute exp(x-max). + const uint DATA_CACHE_SIZE = 16; + FLOAT_TYPE data_cache[DATA_CACHE_SIZE]; + + [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { + const uint col = col0 + tid; + + FLOAT_TYPE a = FLOAT_TYPE(0); + if (col < p.KX) { + a = data_a[rowx * p.KX + col]; + } + + FLOAT_TYPE b = FLOAT_TYPE(0); + if (p.KY > 0 && col < p.KX) { + b = data_b[rowy * p.KX + col]; + } + + FLOAT_TYPE v = a * p.scale + slope * b; + + if (col < p.KX) { + max_val = max(max_val, v); + } + + if (idx < DATA_CACHE_SIZE) { + data_cache[idx] = v; + } + } + + // reduce across the workgroup + vals[tid] = max_val; + barrier(); + [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + vals[tid] = max(vals[tid], vals[tid + s]); + } + barrier(); + } + + max_val = vals[0]; + barrier(); + + FLOAT_TYPE sum = FLOAT_TYPE(0.0f); + + // Compute sum{exp(x - max)} + [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { + const uint col = col0 + tid; + + if (col >= p.KX) { + break; + } + + // compute exp(a*scale+b*slope), add it to sum, and cache the new value + // in data_cache if possible. + const uint i = rowx * p.KX + col; + FLOAT_TYPE val; + if (idx < DATA_CACHE_SIZE) { + val = exp(data_cache[idx] - max_val); + } else { + val = exp(FLOAT_TYPE(data_a[i]) * p.scale + (p.KY > 0 ? slope * FLOAT_TYPE(data_b[rowy * p.KX + col]) : FLOAT_TYPE(0.0f)) - max_val); + } + sum += val; + if (idx < DATA_CACHE_SIZE) { + data_cache[idx] = val; + } else { + data_d[i] = D_TYPE(val); + } + } + + // reduce across the workgroup + vals[tid] = sum; + barrier(); + [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + vals[tid] += vals[tid + s]; + } + barrier(); + } + sum = vals[0]; + + FLOAT_TYPE rcpdivisor = 1.0/sum; + + [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { + const uint col = col0 + tid; + + if (col >= p.KX) { + continue; + } + + if (idx < DATA_CACHE_SIZE) { + data_d[rowx*p.KX + col] = D_TYPE(data_cache[idx] * rcpdivisor); + } else { + data_d[rowx*p.KX + col] *= D_TYPE(rcpdivisor); + } + } +} + +void main() { + // instantiate the soft_max function for several different + // dimensions, to allow loop unrolling + uint num_blocks = (p.KX + BLOCK_SIZE - 1) / BLOCK_SIZE; + if (num_blocks > 32) { + soft_max(num_blocks); + } else if (num_blocks > 16) { + soft_max(32); + } else if (num_blocks > 8) { + soft_max(16); + } else if (num_blocks > 4) { + soft_max(8); + } else if (num_blocks == 4) { + soft_max(4); + } else if (num_blocks == 3) { + soft_max(3); + } else if (num_blocks == 2) { + soft_max(2); + } else if (num_blocks == 1) { + soft_max(1); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp new file mode 100644 index 000000000..ef43598ba --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp @@ -0,0 +1,17 @@ +#version 450 + +#include "types.comp" +#include "generic_unary_head.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + + const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(val * val); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp new file mode 100644 index 000000000..961e5ffa1 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp @@ -0,0 +1,37 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +layout (constant_id = 0) const uint BLOCK_SIZE = 32; + +shared FLOAT_TYPE tmp[BLOCK_SIZE]; + +void main() { + const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; + const uint col = gl_LocalInvocationID.x; + + tmp[col] = FLOAT_TYPE(0.0f); + + for (uint i = col; i < p.KX; i += BLOCK_SIZE) { + tmp[col] += FLOAT_TYPE(data_a[row*p.KX + i]); + } + + barrier(); + [[unroll]] for (int s = int(BLOCK_SIZE) / 2; s > 0; s >>= 1) { + if (col < s) { + tmp[col] += tmp[col + s]; + } + barrier(); + } + + if (col == 0) { + data_d[row] = D_TYPE(tmp[0]); + } +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp new file mode 100644 index 000000000..495f966bd --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp @@ -0,0 +1,20 @@ +#version 450 + +#include "generic_head.comp" +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + data_d[i] = D_TYPE(1. - 2. / (exp(2.*data_a[i]) + 1.)); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp new file mode 100644 index 000000000..28eb24e11 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp @@ -0,0 +1,7 @@ +#version 460 + +#extension GL_NV_cooperative_matrix2 : require + +void main() +{ +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp new file mode 100644 index 000000000..79e065a93 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp @@ -0,0 +1,41 @@ +#version 450 + +#extension GL_EXT_shader_16bit_storage : require + +layout (push_constant) uniform parameter +{ + uint nb1; + uint dim; + uint max_period; +} p; + +#include "types.comp" + +#extension GL_EXT_control_flow_attributes : enable +#define BLOCK_SIZE 256 + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint i = gl_WorkGroupID.y; + const uint j = gl_GlobalInvocationID.x; + const uint d_offset = i * p.nb1; + + if (p.dim % 2 != 0 && j == ((p.dim + 1) / 2)) { + data_d[d_offset + p.dim] = 0.f; + } + + const uint half_dim = p.dim / 2; + if (j >= half_dim) { + return; + } + + const float timestep = float(data_a[i]); + const float freq = float(exp(-log(p.max_period) * j / half_dim)); + const float arg = timestep * freq; + data_d[d_offset + j] = D_TYPE(cos(arg)); + data_d[d_offset + j + half_dim] = D_TYPE(sin(arg)); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp new file mode 100644 index 000000000..eecc47f3a --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp @@ -0,0 +1,323 @@ + +#if !defined(GGML_TYPES_COMP) +#define GGML_TYPES_COMP + +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#if defined(DATA_A_F32) +#define QUANT_K 1 +#define QUANT_R 1 + +#if !defined(LOAD_VEC_A) || LOAD_VEC_A == 1 +#define A_TYPE float +#elif LOAD_VEC_A == 4 +#define A_TYPE vec4 +#elif LOAD_VEC_A == 8 +#define A_TYPE mat2x4 +#endif +#endif + +#if defined(DATA_A_F16) +#define QUANT_K 1 +#define QUANT_R 1 + +#if !defined(LOAD_VEC_A) || LOAD_VEC_A == 1 +#define A_TYPE float16_t +#elif LOAD_VEC_A == 4 +#define A_TYPE f16vec4 +#elif LOAD_VEC_A == 8 +#define A_TYPE f16mat2x4 +#endif +#endif + +#define QUANT_K_Q4_0 32 +#define QUANT_R_Q4_0 2 + +struct block_q4_0 +{ + float16_t d; + uint8_t qs[16]; +}; +struct block_q4_0_packed16 +{ + float16_t d; + uint16_t qs[16/2]; +}; + +#if defined(DATA_A_Q4_0) +#define QUANT_K QUANT_K_Q4_0 +#define QUANT_R QUANT_R_Q4_0 +#define A_TYPE block_q4_0 +#define A_TYPE_PACKED16 block_q4_0_packed16 +#endif + +#define QUANT_K_Q4_1 32 +#define QUANT_R_Q4_1 2 + +struct block_q4_1 +{ + float16_t d; + float16_t m; + uint8_t qs[16]; +}; + +struct block_q4_1_packed16 +{ + float16_t d; + float16_t m; + uint16_t qs[16/2]; +}; + +#if defined(DATA_A_Q4_1) +#define QUANT_K QUANT_K_Q4_1 +#define QUANT_R QUANT_R_Q4_1 +#define A_TYPE block_q4_1 +#define A_TYPE_PACKED16 block_q4_1_packed16 +#endif + +#define QUANT_K_Q5_0 32 +#define QUANT_R_Q5_0 2 + +struct block_q5_0 +{ + float16_t d; + uint16_t qh[2]; + uint8_t qs[16]; +}; + +struct block_q5_0_packed16 +{ + float16_t d; + uint16_t qh[2]; + uint16_t qs[16/2]; +}; + +#if defined(DATA_A_Q5_0) +#define QUANT_K QUANT_K_Q5_0 +#define QUANT_R QUANT_R_Q5_0 +#define A_TYPE block_q5_0 +#define A_TYPE_PACKED16 block_q5_0_packed16 +#endif + +#define QUANT_K_Q5_1 32 +#define QUANT_R_Q5_1 2 + +struct block_q5_1 +{ + float16_t d; + float16_t m; + uint qh; + uint8_t qs[16]; +}; + +struct block_q5_1_packed16 +{ + float16_t d; + float16_t m; + uint qh; + uint16_t qs[16/2]; +}; + +#if defined(DATA_A_Q5_1) +#define QUANT_K QUANT_K_Q5_1 +#define QUANT_R QUANT_R_Q5_1 +#define A_TYPE block_q5_1 +#define A_TYPE_PACKED16 block_q5_1_packed16 +#endif + +#define QUANT_K_Q8_0 32 +#define QUANT_R_Q8_0 1 + +struct block_q8_0 +{ + float16_t d; + int8_t qs[32]; +}; +struct block_q8_0_packed16 +{ + float16_t d; + uint16_t qs[32/2]; +}; + +#if defined(DATA_A_Q8_0) +#define QUANT_K QUANT_K_Q8_0 +#define QUANT_R QUANT_R_Q8_0 +#define A_TYPE block_q8_0 +#define A_TYPE_PACKED16 block_q8_0_packed16 +#endif + +// K-quants +#define QUANT_K_Q2_K 256 + +struct block_q2_K +{ + uint8_t scales[QUANT_K_Q2_K/16]; + uint8_t qs[QUANT_K_Q2_K/4]; + f16vec2 d; +}; + +struct block_q2_K_packed16 +{ + uint16_t scales[QUANT_K_Q2_K/16/2]; + uint16_t qs[QUANT_K_Q2_K/4/2]; + f16vec2 d; +}; + +struct block_q2_K_packed32 +{ + uint32_t scales[QUANT_K_Q2_K/16/4]; + uint32_t qs[QUANT_K_Q2_K/4/4]; + f16vec2 d; +}; + +#if defined(DATA_A_Q2_K) +#define QUANT_K QUANT_K_Q2_K +#define A_TYPE block_q2_K +#define A_TYPE_PACKED16 block_q2_K_packed16 +#define A_TYPE_PACKED32 block_q2_K_packed32 +#endif + +#define QUANT_K_Q3_K 256 + +struct block_q3_K +{ + uint8_t hmask[QUANT_K_Q3_K/8]; + uint8_t qs[QUANT_K_Q3_K/4]; + uint8_t scales[12]; + float16_t d; +}; + +struct block_q3_K_packed16 +{ + uint16_t hmask[QUANT_K_Q3_K/8/2]; + uint16_t qs[QUANT_K_Q3_K/4/2]; + uint16_t scales[12/2]; + float16_t d; +}; + +#if defined(DATA_A_Q3_K) +#define QUANT_K QUANT_K_Q3_K +#define A_TYPE block_q3_K +#define A_TYPE_PACKED16 block_q3_K_packed16 +#endif + +#define QUANT_K_Q4_K 256 + +struct block_q4_K +{ + f16vec2 d; + uint8_t scales[3*QUANT_K_Q4_K/64]; + uint8_t qs[QUANT_K_Q4_K/2]; +}; + +struct block_q4_K_packed16 +{ + f16vec2 d; + uint16_t scales[3*QUANT_K_Q4_K/64/2]; + uint16_t qs[QUANT_K_Q4_K/2/2]; +}; + +struct block_q4_K_packed32 +{ + f16vec2 d; + uint32_t scales[3*QUANT_K_Q4_K/64/4]; + uint32_t qs[QUANT_K_Q4_K/2/4]; +}; + +#if defined(DATA_A_Q4_K) +#define QUANT_K QUANT_K_Q4_K +#define A_TYPE block_q4_K +#define A_TYPE_PACKED16 block_q4_K_packed16 +#define A_TYPE_PACKED32 block_q4_K_packed32 +#endif + +#define QUANT_K_Q5_K 256 + +struct block_q5_K +{ + f16vec2 d; + uint8_t scales[12]; + uint8_t qh[QUANT_K_Q5_K/8]; + uint8_t qs[QUANT_K_Q5_K/2]; +}; + +struct block_q5_K_packed16 +{ + f16vec2 d; + uint16_t scales[12/2]; + uint16_t qh[QUANT_K_Q5_K/8/2]; + uint16_t qs[QUANT_K_Q5_K/2/2]; +}; + +#if defined(DATA_A_Q5_K) +#define QUANT_K QUANT_K_Q5_K +#define A_TYPE block_q5_K +#define A_TYPE_PACKED16 block_q5_K_packed16 +#endif + +#define QUANT_K_Q6_K 256 + +struct block_q6_K +{ + uint8_t ql[QUANT_K_Q6_K/2]; + uint8_t qh[QUANT_K_Q6_K/4]; + int8_t scales[QUANT_K_Q6_K/16]; + float16_t d; +}; + +struct block_q6_K_packed16 +{ + uint16_t ql[QUANT_K_Q6_K/2/2]; + uint16_t qh[QUANT_K_Q6_K/4/2]; + int8_t scales[QUANT_K_Q6_K/16]; + float16_t d; +}; + +#if defined(DATA_A_Q6_K) +#define QUANT_K QUANT_K_Q6_K +#define A_TYPE block_q6_K +#define A_TYPE_PACKED16 block_q6_K_packed16 +#endif + +// IQuants + +#define QUANT_K_IQ4_NL 32 +#define QUANT_R_IQ4_NL 2 + +struct block_iq4_nl +{ + float16_t d; + uint8_t qs[QUANT_K_IQ4_NL/2]; +}; + +struct block_iq4_nl_packed16 +{ + float16_t d; + uint16_t qs[QUANT_K_IQ4_NL/2/2]; +}; + +#if defined(DATA_A_IQ4_NL) + +const int8_t kvalues_iq4nl_const[16] = { + int8_t(-127), int8_t(-104), int8_t(-83), int8_t(-65), int8_t(-49), int8_t(-35), int8_t(-22), int8_t(-10), + int8_t(1), int8_t(13), int8_t(25), int8_t(38), int8_t(53), int8_t(69), int8_t(89), int8_t(113) +}; + +shared FLOAT_TYPE kvalues_iq4nl[16]; + +void init_iq4nl_shmem() +{ + // copy the table into shared memory and sync + if (gl_LocalInvocationIndex.x < 16) { + kvalues_iq4nl[gl_LocalInvocationIndex.x] = FLOAT_TYPE(kvalues_iq4nl_const[gl_LocalInvocationIndex.x]); + } + barrier(); +} + +#define QUANT_K QUANT_K_IQ4_NL +#define QUANT_R QUANT_R_IQ4_NL +#define A_TYPE block_iq4_nl +#define A_TYPE_PACKED16 block_iq4_nl_packed16 +#endif + +#endif // !defined(GGML_TYPES_COMP) diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp new file mode 100644 index 000000000..6f607380d --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp @@ -0,0 +1,36 @@ +#version 450 + +layout (push_constant) uniform parameter +{ + uint ne; uint a_offset; uint d_offset; + uint nb00; uint nb01; uint nb02; uint nb03; + uint ne10; uint ne11; uint ne12; uint ne13; + float sf0; float sf1; float sf2; float sf3; +} p; + +#include "types.comp" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (idx >= p.ne) { + return; + } + + const uint i10 = idx % p.ne10; + const uint i11 = (idx / p.ne10) % p.ne11; + const uint i12 = (idx / (p.ne10 * p.ne11)) % p.ne12; + const uint i13 = (idx / (p.ne10 * p.ne11 * p.ne12)) % p.ne13; + + const uint i00 = uint(i10 / p.sf0); + const uint i01 = uint(i11 / p.sf1); + const uint i02 = uint(i12 / p.sf2); + const uint i03 = uint(i13 / p.sf3); + + data_d[p.d_offset + idx] = D_TYPE(data_a[p.a_offset + i03 * p.nb03 + i02 * p.nb02 + i01 * p.nb01 + i00 * p.nb00]); +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp new file mode 100644 index 000000000..8111c0638 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -0,0 +1,594 @@ + + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef _WIN32 + #include + #include // For _mkdir on Windows + #include // For std::replace on w64devkit +#else + #include + #include + #include +#endif + +#include + +#define ASYNCIO_CONCURRENCY 64 + +std::mutex lock; +std::vector> shader_fnames; + +std::string GLSLC = "glslc"; +std::string input_dir = "vulkan-shaders"; +std::string output_dir = "/tmp"; +std::string target_hpp = "ggml-vulkan-shaders.hpp"; +std::string target_cpp = "ggml-vulkan-shaders.cpp"; +bool no_clean = false; + +const std::vector type_names = { + "f32", + "f16", + "q4_0", + "q4_1", + "q5_0", + "q5_1", + "q8_0", + "q2_k", + "q3_k", + "q4_k", + "q5_k", + "q6_k", + "iq4_nl" +}; + +namespace { +void execute_command(const std::string& command, std::string& stdout_str, std::string& stderr_str) { +#ifdef _WIN32 + HANDLE stdout_read, stdout_write; + HANDLE stderr_read, stderr_write; + SECURITY_ATTRIBUTES sa = { sizeof(SECURITY_ATTRIBUTES), NULL, TRUE }; + + if (!CreatePipe(&stdout_read, &stdout_write, &sa, 0) || + !SetHandleInformation(stdout_read, HANDLE_FLAG_INHERIT, 0)) { + throw std::runtime_error("Failed to create stdout pipe"); + } + + if (!CreatePipe(&stderr_read, &stderr_write, &sa, 0) || + !SetHandleInformation(stderr_read, HANDLE_FLAG_INHERIT, 0)) { + throw std::runtime_error("Failed to create stderr pipe"); + } + + PROCESS_INFORMATION pi; + STARTUPINFOA si = {}; + si.cb = sizeof(STARTUPINFOA); + si.dwFlags = STARTF_USESTDHANDLES; + si.hStdOutput = stdout_write; + si.hStdError = stderr_write; + + std::vector cmd(command.begin(), command.end()); + cmd.push_back('\0'); + + if (!CreateProcessA(NULL, cmd.data(), NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi)) { + throw std::runtime_error("Failed to create process"); + } + + CloseHandle(stdout_write); + CloseHandle(stderr_write); + + std::array buffer; + DWORD bytes_read; + + while (ReadFile(stdout_read, buffer.data(), (DWORD)buffer.size(), &bytes_read, NULL) && bytes_read > 0) { + stdout_str.append(buffer.data(), bytes_read); + } + + while (ReadFile(stderr_read, buffer.data(), (DWORD)buffer.size(), &bytes_read, NULL) && bytes_read > 0) { + stderr_str.append(buffer.data(), bytes_read); + } + + CloseHandle(stdout_read); + CloseHandle(stderr_read); + WaitForSingleObject(pi.hProcess, INFINITE); + CloseHandle(pi.hProcess); + CloseHandle(pi.hThread); +#else +int stdout_pipe[2]; + int stderr_pipe[2]; + + if (pipe(stdout_pipe) != 0 || pipe(stderr_pipe) != 0) { + throw std::runtime_error("Failed to create pipes"); + } + + pid_t pid = fork(); + if (pid < 0) { + throw std::runtime_error("Failed to fork process"); + } + + if (pid == 0) { + close(stdout_pipe[0]); + close(stderr_pipe[0]); + dup2(stdout_pipe[1], STDOUT_FILENO); + dup2(stderr_pipe[1], STDERR_FILENO); + close(stdout_pipe[1]); + close(stderr_pipe[1]); + execl("/bin/sh", "sh", "-c", command.c_str(), (char*) nullptr); + _exit(EXIT_FAILURE); + } else { + close(stdout_pipe[1]); + close(stderr_pipe[1]); + + std::array buffer; + ssize_t bytes_read; + + while ((bytes_read = read(stdout_pipe[0], buffer.data(), buffer.size())) > 0) { + stdout_str.append(buffer.data(), bytes_read); + } + + while ((bytes_read = read(stderr_pipe[0], buffer.data(), buffer.size())) > 0) { + stderr_str.append(buffer.data(), bytes_read); + } + + close(stdout_pipe[0]); + close(stderr_pipe[0]); + waitpid(pid, nullptr, 0); + } +#endif +} + +bool directory_exists(const std::string& path) { + struct stat info; + if (stat(path.c_str(), &info) != 0) { + return false; // Path doesn't exist or can't be accessed + } + return (info.st_mode & S_IFDIR) != 0; // Check if it is a directory +} + +bool create_directory(const std::string& path) { +#ifdef _WIN32 + return _mkdir(path.c_str()) == 0 || errno == EEXIST; // EEXIST means the directory already exists +#else + return mkdir(path.c_str(), 0755) == 0 || errno == EEXIST; // 0755 is the directory permissions +#endif +} + +std::string to_uppercase(const std::string& input) { + std::string result = input; + for (char& c : result) { + c = std::toupper(c); + } + return result; +} + +bool string_ends_with(const std::string& str, const std::string& suffix) { + if (suffix.size() > str.size()) { + return false; + } + return std::equal(suffix.rbegin(), suffix.rend(), str.rbegin()); +} + +static const char path_separator = '/'; + +std::string join_paths(const std::string& path1, const std::string& path2) { + return path1 + path_separator + path2; +} + +std::string basename(const std::string &path) { + return path.substr(path.find_last_of("/\\") + 1); +} + +// variables to track number of compiles in progress +static uint32_t compile_count = 0; +static std::mutex compile_count_mutex; +static std::condition_variable compile_count_cond; + +void string_to_spv_func(const std::string& _name, const std::string& in_fname, const std::map& defines, bool fp16 = true, bool coopmat = false, bool coopmat2 = false, bool f16acc = false) { + std::string name = _name + (f16acc ? "_f16acc" : "") + (coopmat ? "_coopmat" : "") + (coopmat2 ? "_cm2" : (fp16 ? "" : "_fp32")); + std::string out_fname = join_paths(output_dir, name + ".spv"); + std::string in_path = join_paths(input_dir, in_fname); + + std::string target_env = (name.find("_cm2") != std::string::npos) ? "--target-env=vulkan1.3" : "--target-env=vulkan1.2"; + + // disable spirv-opt for coopmat shaders for https://github.com/ggerganov/llama.cpp/issues/10734 + std::string opt_level = coopmat ? "" : "-O"; + + #ifdef _WIN32 + std::vector cmd = {GLSLC, "-fshader-stage=compute", target_env, opt_level, "\"" + in_path + "\"", "-o", "\"" + out_fname + "\""}; + #else + std::vector cmd = {GLSLC, "-fshader-stage=compute", target_env, opt_level, in_path, "-o", out_fname}; + #endif + + #ifdef GGML_VULKAN_SHADER_DEBUG_INFO + cmd.push_back("-g"); + #endif + + for (const auto& define : defines) { + cmd.push_back("-D" + define.first + "=" + define.second); + } + + std::string command; + for (const auto& part : cmd) { + command += part + " "; + } + + std::string stdout_str, stderr_str; + try { + // std::cout << "Executing command: "; + // for (const auto& part : cmd) { + // std::cout << part << " "; + // } + // std::cout << std::endl; + + execute_command(command, stdout_str, stderr_str); + if (!stderr_str.empty()) { + std::cerr << "cannot compile " << name << "\n\n" << command << "\n\n" << stderr_str << std::endl; + return; + } + + std::lock_guard guard(lock); + shader_fnames.push_back(std::make_pair(name, out_fname)); + } catch (const std::exception& e) { + std::cerr << "Error executing command for " << name << ": " << e.what() << std::endl; + } + { + std::lock_guard guard(compile_count_mutex); + assert(compile_count > 0); + compile_count--; + } + compile_count_cond.notify_all(); +} + +std::map merge_maps(const std::map& a, const std::map& b) { + std::map result = a; + result.insert(b.begin(), b.end()); + return result; +} + +static std::vector> compiles; +void string_to_spv(const std::string& _name, const std::string& in_fname, const std::map& defines, bool fp16 = true, bool coopmat = false, bool coopmat2 = false, bool f16acc = false) { + { + // wait until fewer than N compiles are in progress. + // 16 is an arbitrary limit, the goal is to avoid "failed to create pipe" errors. + uint32_t N = 16; + std::unique_lock guard(compile_count_mutex); + while (compile_count >= N) { + compile_count_cond.wait(guard); + } + compile_count++; + } + compiles.push_back(std::async(string_to_spv_func, _name, in_fname, defines, fp16, coopmat, coopmat2, f16acc)); +} + +void matmul_shaders(bool fp16, bool matmul_id, bool coopmat, bool coopmat2, bool f16acc) { + std::string load_vec = coopmat2 ? "1" : fp16 ? "8" : "4"; + std::string aligned_b_type_f32 = coopmat2 ? "float" : fp16 ? "mat2x4" : "vec4"; + std::string aligned_b_type_f16 = coopmat2 ? "float16_t" : fp16 ? "f16mat2x4" : "f16vec4"; + + std::map base_dict = {{"FLOAT_TYPE", (coopmat2 || fp16) ? "float16_t" : "float"}}; + std::string shader_name = "matmul"; + + if (matmul_id) { + base_dict["MUL_MAT_ID"] = "1"; + shader_name = "matmul_id"; + } + + if (fp16) { + base_dict["FLOAT16"] = "1"; + } + + base_dict["ACC_TYPE"] = f16acc ? "float16_t" : "float"; + + if (coopmat) { + base_dict["COOPMAT"] = "1"; + } + + base_dict["ACC_TYPE"] = f16acc ? "float16_t" : "float"; + + std::string source_name = coopmat2 ? "mul_mm_cm2.comp" : "mul_mm.comp"; + + // Shaders with f16 B_TYPE + string_to_spv(shader_name + "_f32_f16", source_name, merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}, }), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_f32_f16_aligned", source_name, merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + + string_to_spv(shader_name + "_f16_aligned", source_name, merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_f16", source_name, merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); + + for (const auto& tname : type_names) { + std::string data_a_key = "DATA_A_" + to_uppercase(tname); + // For unaligned, load one at a time for f32/f16, or two at a time for quants + std::string load_vec_a_unaligned = (coopmat2 || tname == "f32" || tname == "f16") ? "1" : "2"; + // For aligned matmul loads + std::string load_vec_a = (coopmat2 || tname == "f32" || tname == "f16") ? load_vec : "2"; + + string_to_spv(shader_name + "_" + tname + "_f32", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_" + tname + "_f32_aligned", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + + if (tname != "f16" && tname != "f32") { + string_to_spv(shader_name + "_" + tname + "_f16", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_" + tname + "_f16_aligned", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + } + } +} + +void process_shaders() { + std::cout << "ggml_vulkan: Generating and compiling shaders to SPIR-V" << std::endl; + std::map base_dict = {{"FLOAT_TYPE", "float"}}; + + // matmul + for (const auto& matmul_id : {false, true}) { + // No coopmats + // fp32 + matmul_shaders(false, matmul_id, false, false, false); + + // fp16, fp32acc and fp16acc + matmul_shaders(true, matmul_id, false, false, false); + matmul_shaders(true, matmul_id, false, false, true); + + // Coopmat, fp32acc and fp16acc + matmul_shaders(true, matmul_id, true, false, false); + matmul_shaders(true, matmul_id, true, false, true); + +#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + // Coopmat2, fp32acc and fp16acc + matmul_shaders(true, matmul_id, false, true, false); + matmul_shaders(true, matmul_id, false, true, true); +#endif + } + +#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) + // flash attention + for (const auto& f16acc : {false, true}) { + std::string acctype = f16acc ? "float16_t" : "float"; + + for (const auto& tname : type_names) { + if (tname == "f32") { + continue; + } + + if (tname == "f16") { + string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm2.comp", + merge_maps(base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"ACC_TYPE", acctype}}), true, false, true, f16acc); + } else { + std::string data_a_key = "DATA_A_" + to_uppercase(tname); + string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm2.comp", + merge_maps(base_dict, {{data_a_key, "1"}, {"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"ACC_TYPE", acctype}, {"DEQUANTFUNC", "dequantFunc"+to_uppercase(tname) }, {"BLOCK_SIZE", "QUANT_K_"+to_uppercase(tname) }}), true, false, true, f16acc); + } + } + } +#endif + + for (const auto& tname : type_names) { + // mul mat vec + std::string data_a_key = "DATA_A_" + to_uppercase(tname); + std::string shader = (string_ends_with(tname, "_k")) ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp"; + + string_to_spv("mul_mat_vec_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPE_VEC2", "vec2"}, {"B_TYPE_VEC4", "vec4"}, {"D_TYPE", "float"}})); + string_to_spv("mul_mat_vec_" + tname + "_f16_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"B_TYPE_VEC2", "f16vec2"}, {"B_TYPE_VEC4", "f16vec4"}, {"D_TYPE", "float"}})); + + string_to_spv("mul_mat_vec_id_" + tname + "_f32", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPE_VEC2", "vec2"}, {"B_TYPE_VEC4", "vec4"}, {"D_TYPE", "float"}})); + + // Dequant shaders + if (tname != "f16") { + string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}})); + } + + if (!string_ends_with(tname, "_k")) { + shader = (tname == "f32" || tname == "f16") ? "get_rows.comp" : "get_rows_quant.comp"; + + if (tname == "f16") { + string_to_spv("get_rows_" + tname, shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}})); + } else { + string_to_spv("get_rows_" + tname, shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}})); + } + string_to_spv("get_rows_" + tname + "_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float"}})); + } + } + + string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("mul_mat_vec_nc_f16_f32", "mul_mat_vec_nc.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); + + // Norms + string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + + string_to_spv("cpy_f32_f32", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("cpy_f32_f16", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); + string_to_spv("cpy_f16_f16", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); + string_to_spv("contig_cpy_f32_f32", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("contig_cpy_f32_f16", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); + string_to_spv("contig_cpy_f16_f16", "contig_copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); + + string_to_spv("add_f32", "add.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + string_to_spv("add_f16_f32_f16", "add.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("acc_f32", "acc.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("split_k_reduce", "mul_mat_split_k_reduce.comp", {}); + + string_to_spv("mul_f32", "mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("div_f32", "div.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("repeat_f32", "repeat.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + + string_to_spv("scale_f32", "scale.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("sqr_f32", "square.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("sin_f32", "sin.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("cos_f32", "cos.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("clamp_f32", "clamp.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + + string_to_spv("pad_f32", "pad.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + + string_to_spv("concat_f32", "concat.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("concat_f16", "concat.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); + string_to_spv("concat_i32", "concat.comp", {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}}); + + string_to_spv("upscale_f32", "upscale.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); + + string_to_spv("gelu_f32", "gelu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("gelu_quick_f32", "gelu_quick.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("silu_f32", "silu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("relu_f32", "relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("leaky_relu_f32", "leaky_relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("tanh_f32", "tanh.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + + string_to_spv("diag_mask_inf_f32", "diag_mask_inf.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + + string_to_spv("soft_max_f32", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("soft_max_f32_f16", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}})); + + string_to_spv("rope_norm_f32", "rope_norm.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("rope_norm_f16", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); + string_to_spv("rope_norm_f16_rte", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}}); + + string_to_spv("rope_neox_f32", "rope_neox.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("rope_neox_f16", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); + string_to_spv("rope_neox_f16_rte", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}}); + + string_to_spv("argsort_f32", "argsort.comp", {{"A_TYPE", "float"}}); + + string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + + string_to_spv("im2col_f32", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("im2col_f32_f16", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}})); + string_to_spv("im2col_f32_f16_rte", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}})); + + string_to_spv("timestep_embedding_f32", "timestep_embedding.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + + string_to_spv("pool2d_f32", "pool2d.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + + string_to_spv("rwkv_wkv6_f32", "wkv6.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); + + for (auto &c : compiles) { + c.wait(); + } +} + +void write_output_files() { + FILE* hdr = fopen(target_hpp.c_str(), "w"); + FILE* src = fopen(target_cpp.c_str(), "w"); + + fprintf(hdr, "#include \n\n"); + fprintf(src, "#include \"%s\"\n\n", basename(target_hpp).c_str()); + + for (const auto& pair : shader_fnames) { + const std::string& name = pair.first; + #ifdef _WIN32 + std::string path = pair.second; + std::replace(path.begin(), path.end(), '/', '\\' ); + #else + const std::string& path = pair.second; + #endif + + FILE* spv = fopen(path.c_str(), "rb"); + if (!spv) { + std::cerr << "Error opening SPIR-V file: " << path << " (" << strerror(errno) << ")\n"; + continue; + } + + fseek(spv, 0, SEEK_END); + size_t size = ftell(spv); + fseek(spv, 0, SEEK_SET); + + std::vector data(size); + size_t read_size = fread(data.data(), 1, size, spv); + fclose(spv); + if (read_size != size) { + std::cerr << "Error reading SPIR-V file: " << path << " (" << strerror(errno) << ")\n"; + continue; + } + + fprintf(hdr, "extern unsigned char %s_data[%zu];\n", name.c_str(), size); + fprintf(hdr, "const uint64_t %s_len = %zu;\n\n", name.c_str(), size); + + fprintf(src, "unsigned char %s_data[%zu] = {\n", name.c_str(), size); + for (size_t i = 0; i < size; ++i) { + fprintf(src, "0x%02x,", data[i]); + if ((i + 1) % 12 == 0) fprintf(src, "\n"); + } + fprintf(src, "\n};\n\n"); + + if (!no_clean) { + std::remove(path.c_str()); + } + } + + fclose(hdr); + fclose(src); +} +} + +int main(int argc, char** argv) { + std::map args; + for (int i = 1; i < argc; ++i) { + std::string arg = argv[i]; + if (arg.rfind("--", 0) == 0) { + if (i + 1 < argc && argv[i + 1][0] != '-') { + args[arg] = argv[i + 1]; + ++i; + } else { + args[arg] = ""; + } + } + } + + if (args.find("--glslc") != args.end()) { + GLSLC = args["--glslc"]; // Path to glslc + } + if (args.find("--input-dir") != args.end()) { + input_dir = args["--input-dir"]; // Directory containing shader sources + } + if (args.find("--output-dir") != args.end()) { + output_dir = args["--output-dir"]; // Directory for containing SPIR-V output + } + if (args.find("--target-hpp") != args.end()) { + target_hpp = args["--target-hpp"]; // Path to generated header file + } + if (args.find("--target-cpp") != args.end()) { + target_cpp = args["--target-cpp"]; // Path to generated cpp file + } + if (args.find("--no-clean") != args.end()) { + no_clean = true; // Keep temporary SPIR-V files in output-dir after build + } + + if (!directory_exists(input_dir)) { + std::cerr << "\"" << input_dir << "\" must be a valid directory containing shader sources" << std::endl; + return EXIT_FAILURE; + } + + if (!directory_exists(output_dir)) { + if (!create_directory(output_dir)) { + std::cerr << "Error creating output directory: " << output_dir << "\n"; + return EXIT_FAILURE; + } + } + + process_shaders(); + + write_output_files(); + + return EXIT_SUCCESS; +} diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp new file mode 100644 index 000000000..35cc6c45f --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp @@ -0,0 +1,87 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +#define BLOCK_SIZE 64 +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform Parameters { + uint B; + uint T; + uint C; + uint H; +}; + +layout(binding = 0) readonly buffer KBuf { A_TYPE k[]; }; +layout(binding = 1) readonly buffer VBuf { A_TYPE v[]; }; +layout(binding = 2) readonly buffer RBuf { A_TYPE r[]; }; +layout(binding = 3) readonly buffer TimeFBuf { A_TYPE tf[]; }; +layout(binding = 4) readonly buffer TimeDBuf { A_TYPE td[]; }; +layout(binding = 5) readonly buffer StateBuf { A_TYPE state_in[]; }; +layout(binding = 6) buffer DstBuf { A_TYPE dst[]; }; + +shared A_TYPE _k[BLOCK_SIZE], _r[BLOCK_SIZE], _tf[BLOCK_SIZE], _td[BLOCK_SIZE]; + +void main() { + const uint head_size = BLOCK_SIZE; + const uint batch_id = gl_WorkGroupID.x / H; + const uint head_id = gl_WorkGroupID.x % H; + const uint tid = gl_LocalInvocationID.x; + + const uint state_size = C * head_size; + const uint n_seq_tokens = T / B; + + if (batch_id >= B || head_id >= H) { + return; + } + + A_TYPE state[BLOCK_SIZE]; + [[unroll]] for (uint i = 0; i < head_size; i++) { + state[i] = state_in[batch_id * state_size + head_id * head_size * head_size + + i * head_size + tid]; + } + + barrier(); + _tf[tid] = tf[head_id * head_size + tid]; + barrier(); + + const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid; + const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid; + + for (uint t = start_t; t < end_t; t += C) { + barrier(); + _k[tid] = k[t]; + _r[tid] = r[t]; + _td[tid] = td[t]; + barrier(); + + const A_TYPE v_val = v[t]; + A_TYPE y = 0.0; + + [[unroll]] for (uint j = 0; j < head_size; j += 4) { + vec4 k_vec = vec4(_k[j], _k[j+1], _k[j+2], _k[j+3]); + vec4 r_vec = vec4(_r[j], _r[j+1], _r[j+2], _r[j+3]); + vec4 tf_vec = vec4(_tf[j], _tf[j+1], _tf[j+2], _tf[j+3]); + vec4 td_vec = vec4(_td[j], _td[j+1], _td[j+2], _td[j+3]); + vec4 s_vec = vec4(state[j], state[j+1], state[j+2], state[j+3]); + + vec4 kv = k_vec * v_val; + + vec4 temp = tf_vec * kv + s_vec; + y += dot(r_vec, temp); + + s_vec = s_vec * td_vec + kv; + state[j] = s_vec.x; + state[j+1] = s_vec.y; + state[j+2] = s_vec.z; + state[j+3] = s_vec.w; + } + + dst[t] = y; + } + + [[unroll]] for (uint i = 0; i < head_size; i++) { + dst[T * C + batch_id * state_size + head_id * head_size * head_size + + i * head_size + tid] = state[i]; + } +} From b14dd68feeaebafc71479f8df2ef033ccb0bac3d Mon Sep 17 00:00:00 2001 From: Vadim Grinco Date: Mon, 10 Mar 2025 12:15:43 +0100 Subject: [PATCH 26/29] Fixed the "detached head" issues Signed-off-by: Vadim Grinco --- Dockerfile | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/Dockerfile b/Dockerfile index 9e2928108..a31706243 100644 --- a/Dockerfile +++ b/Dockerfile @@ -37,13 +37,12 @@ COPY patches/ /tmp/patches/ RUN \ git clone https://github.com/pufferffish/ollama-vulkan.git "/tmp/ollama-vulkan-git" && \ cd "/tmp/ollama-vulkan-git" && \ - git checkout 2d443b3dd660a1fd2760d64538512df93648b4bb && git checkout -b ollama_vulkan_stable && \ + git checkout 2d443b3dd660a1fd2760d64538512df93648b4bb -b ollama_vulkan_stable && \ git config user.name "Builder" && git config user.email "builder@local" && \ git remote add ollama_vanilla https://github.com/ollama/ollama.git && \ - git fetch ollama_vanilla --tags && git checkout v0.5.14-rc0 && git checkout -b ollama_vanilla_stable && \ + git fetch ollama_vanilla --tags && git checkout v0.5.13 -b ollama_vanilla_stable && \ git checkout ollama_vulkan_stable && git merge ollama_vanilla_stable --allow-unrelated-histories --no-edit && \ - for p in /tmp/patches/00-fix-vulkan-building.patch; do patch -p1 < $p; done - + for p in /tmp/patches/*.patch; do patch -p1 < $p; done RUN \ cd "/tmp/ollama-vulkan-git" && \ make -f Makefile.sync clean sync From 31606b2febe9b5680cf4fb8e20f26a78dc39ee6e Mon Sep 17 00:00:00 2001 From: Vadim Grinco Date: Mon, 10 Mar 2025 12:49:47 +0100 Subject: [PATCH 27/29] Merged in the right direction Signed-off-by: Vadim Grinco --- Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Dockerfile b/Dockerfile index a31706243..cabb984fc 100644 --- a/Dockerfile +++ b/Dockerfile @@ -41,7 +41,7 @@ RUN \ git config user.name "Builder" && git config user.email "builder@local" && \ git remote add ollama_vanilla https://github.com/ollama/ollama.git && \ git fetch ollama_vanilla --tags && git checkout v0.5.13 -b ollama_vanilla_stable && \ - git checkout ollama_vulkan_stable && git merge ollama_vanilla_stable --allow-unrelated-histories --no-edit && \ + git checkout ollama_vanilla_stable && git merge ollama_vulkan_stable --allow-unrelated-histories --no-edit && \ for p in /tmp/patches/*.patch; do patch -p1 < $p; done RUN \ cd "/tmp/ollama-vulkan-git" && \ From 6b1f84e171ea5fbf88f0dfca8d00856d8e4e4555 Mon Sep 17 00:00:00 2001 From: Vadim Grinco Date: Tue, 11 Mar 2025 14:09:47 +0100 Subject: [PATCH 28/29] Merging the latest stable (#2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Applied 00-fix-vulkan-building.patch * Implemented vulkan backend based on the work done by whyvl, Dts0, McBane87 and others Tested on AMD Ryzen 7 8845HS w/ Radeon 780M Graphics with ROCm disabled ``` [GIN-debug] POST /v1/chat/completions --> github.com/ollama/ollama/server.(*Server).ChatHandler-fm (6 handlers) [GIN-debug] POST /v1/completions --> github.com/ollama/ollama/server.(*Server).GenerateHandler-fm (6 handlers) [GIN-debug] POST /v1/embeddings --> github.com/ollama/ollama/server.(*Server).EmbedHandler-fm (6 handlers) [GIN-debug] GET /v1/models --> github.com/ollama/ollama/server.(*Server).ListHandler-fm (6 handlers) [GIN-debug] GET /v1/models/:model --> github.com/ollama/ollama/server.(*Server).ShowHandler-fm (6 handlers) time=2025-03-11T13:00:40.793Z level=INFO source=gpu.go:199 msg="vulkan: load libvulkan and libcap ok" time=2025-03-11T13:00:40.877Z level=INFO source=gpu.go:421 msg="error looking up vulkan GPU memory" error="device is a CPU" time=2025-03-11T13:00:40.878Z level=WARN source=amd_linux.go:443 msg="amdgpu detected, but no compatible rocm library found. Either install rocm v6, or follow manual install instructions at https://github.com/ollama/ollama/blob/main/docs/linux.md#manual-install" time=2025-03-11T13:00:40.878Z level=WARN source=amd_linux.go:348 msg="unable to verify rocm library: no suitable rocm found, falling back to CPU" time=2025-03-11T13:00:40.879Z level=INFO source=types.go:137 msg="inference compute" id=0 library=vulkan variant="" compute=1.3 driver=1.3 name="AMD Radeon Graphics (RADV GFX1103_R1)" total="15.6 GiB" available="15.6 GiB" ``` ``` # ollama run phi4:14b >>> /set verbose Set 'verbose' mode. >>> how's it going? Hello! I'm here to help you with any questions or tasks you have. How can I assist you today? 😊 total duration: 3.341959745s load duration: 18.165612ms prompt eval count: 15 token(s) prompt eval duration: 475ms prompt eval rate: 31.58 tokens/s eval count: 26 token(s) eval duration: 2.846s eval rate: 9.14 tokens/s >>> ``` --- Dockerfile | 217 +++++++++++++++++++++++++++++++++-------------------- 1 file changed, 135 insertions(+), 82 deletions(-) diff --git a/Dockerfile b/Dockerfile index cabb984fc..870adfd72 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,98 +1,151 @@ -FROM --platform=linux/amd64 library/ubuntu:noble as builder +# vim: filetype=dockerfile -ENV DEBIAN_FRONTEND="noninteractive" +ARG FLAVOR=${TARGETARCH} -ENV VULKAN_VER_BASE="1.3.296" -ENV VULKAN_VER="${VULKAN_VER_BASE}.0" -ENV UBUNTU_VERSION="noble" +ARG ROCMVERSION=6.3.3 +ARG JETPACK5VERSION=r35.4.1 +ARG JETPACK6VERSION=r36.4.0 +ARG CMAKEVERSION=3.31.2 +ARG VULKANVERSION=1.4.304.1 -ENV GOLANG_VERSION="1.22.8" -ENV GOARCH="amd64" -ENV CGO_ENABLED=1 +# CUDA v11 requires gcc v10. v10.3 has regressions, so the rockylinux 8.5 AppStream has the latest compatible version +FROM --platform=linux/amd64 rocm/dev-almalinux-8:${ROCMVERSION}-complete AS base-amd64 +RUN yum install -y yum-utils \ + && yum-config-manager --add-repo https://dl.rockylinux.org/vault/rocky/8.5/AppStream/\$basearch/os/ \ + && rpm --import https://dl.rockylinux.org/pub/rocky/RPM-GPG-KEY-Rocky-8 \ + && dnf install -y yum-utils ccache gcc-toolset-10-gcc-10.2.1-8.2.el8 gcc-toolset-10-gcc-c++-10.2.1-8.2.el8 gcc-toolset-10-binutils-2.35-11.el8 \ + && yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo +ENV PATH=/opt/rh/gcc-toolset-10/root/usr/bin:$PATH +ARG VULKANVERSION +RUN wget https://sdk.lunarg.com/sdk/download/${VULKANVERSION}/linux/vulkansdk-linux-x86_64-${VULKANVERSION}.tar.xz -O /tmp/vulkansdk-linux-x86_64-${VULKANVERSION}.tar.xz \ + && tar xvf /tmp/vulkansdk-linux-x86_64-${VULKANVERSION}.tar.xz \ + && dnf -y install ninja-build libcap-devel \ + && ln -s /usr/bin/python3 /usr/bin/python \ + && /${VULKANVERSION}/vulkansdk -j 8 vulkan-headers \ + && /${VULKANVERSION}/vulkansdk -j 8 shaderc +RUN cp -r /${VULKANVERSION}/x86_64/include/* /usr/local/include/ \ + && cp -r /${VULKANVERSION}/x86_64/lib/* /usr/local/lib +ENV PATH=/${VULKANVERSION}/x86_64/bin:$PATH + +FROM --platform=linux/arm64 almalinux:8 AS base-arm64 +# install epel-release for ccache +RUN yum install -y yum-utils epel-release \ + && dnf install -y clang ccache \ + && yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/sbsa/cuda-rhel8.repo +ENV CC=clang CXX=clang++ + +FROM base-${TARGETARCH} AS base +ARG CMAKEVERSION +RUN curl -fsSL https://github.com/Kitware/CMake/releases/download/v${CMAKEVERSION}/cmake-${CMAKEVERSION}-linux-$(uname -m).tar.gz | tar xz -C /usr/local --strip-components 1 +COPY CMakeLists.txt CMakePresets.json . +COPY ml/backend/ggml/ggml ml/backend/ggml/ggml ENV LDFLAGS=-s -# Default mirror was very slow -RUN \ - sed -i 's/archive.ubuntu.com/gb.archive.ubuntu.com/g' /etc/apt/sources.list.d/ubuntu.sources +FROM base AS cpu +RUN dnf install -y gcc-toolset-11-gcc gcc-toolset-11-gcc-c++ +ENV PATH=/opt/rh/gcc-toolset-11/root/usr/bin:$PATH +RUN --mount=type=cache,target=/root/.ccache \ + cmake --preset 'CPU' \ + && cmake --build --parallel --preset 'CPU' \ + && cmake --install build --component CPU --strip --parallel 8 -RUN \ - apt-get update && \ - apt-get install -y ca-certificates build-essential ccache cmake wget git curl rsync xz-utils libcap-dev +FROM base AS cuda-11 +ARG CUDA11VERSION=11.3 +RUN dnf install -y cuda-toolkit-${CUDA11VERSION//./-} +ENV PATH=/usr/local/cuda-11/bin:$PATH +RUN --mount=type=cache,target=/root/.ccache \ + cmake --preset 'CUDA 11' \ + && cmake --build --parallel --preset 'CUDA 11' \ + && cmake --install build --component CUDA --strip --parallel 8 -RUN \ - mkdir -p /usr/local 2>/dev/null || true && \ - curl -s -L https://dl.google.com/go/go${GOLANG_VERSION}.linux-${GOARCH}.tar.gz | tar -xz -C /usr/local && \ - ln -s /usr/local/go/bin/go /usr/local/bin/go && \ - ln -s /usr/local/go/bin/gofmt /usr/local/bin/gofmt +FROM base AS cuda-12 +ARG CUDA12VERSION=12.8 +RUN dnf install -y cuda-toolkit-${CUDA12VERSION//./-} +ENV PATH=/usr/local/cuda-12/bin:$PATH +RUN --mount=type=cache,target=/root/.ccache \ + cmake --preset 'CUDA 12' \ + && cmake --build --parallel --preset 'CUDA 12' \ + && cmake --install build --component CUDA --strip --parallel 8 + +FROM base AS rocm-6 +ENV PATH=/opt/rocm/hcc/bin:/opt/rocm/hip/bin:/opt/rocm/bin:/opt/rocm/hcc/bin:$PATH +RUN --mount=type=cache,target=/root/.ccache \ + cmake --preset 'ROCm 6' \ + && cmake --build --parallel --preset 'ROCm 6' \ + && cmake --install build --component HIP --strip --parallel 8 + +FROM --platform=linux/arm64 nvcr.io/nvidia/l4t-jetpack:${JETPACK5VERSION} AS jetpack-5 +ARG CMAKEVERSION +RUN apt-get update && apt-get install -y curl ccache \ + && curl -fsSL https://github.com/Kitware/CMake/releases/download/v${CMAKEVERSION}/cmake-${CMAKEVERSION}-linux-$(uname -m).tar.gz | tar xz -C /usr/local --strip-components 1 +COPY CMakeLists.txt CMakePresets.json . +COPY ml/backend/ggml/ggml ml/backend/ggml/ggml +RUN --mount=type=cache,target=/root/.ccache \ + cmake --preset 'JetPack 5' \ + && cmake --build --parallel --preset 'JetPack 5' \ + && cmake --install build --component CUDA --strip --parallel 8 + +FROM --platform=linux/arm64 nvcr.io/nvidia/l4t-jetpack:${JETPACK6VERSION} AS jetpack-6 +ARG CMAKEVERSION +RUN apt-get update && apt-get install -y curl ccache \ + && curl -fsSL https://github.com/Kitware/CMake/releases/download/v${CMAKEVERSION}/cmake-${CMAKEVERSION}-linux-$(uname -m).tar.gz | tar xz -C /usr/local --strip-components 1 +COPY CMakeLists.txt CMakePresets.json . +COPY ml/backend/ggml/ggml ml/backend/ggml/ggml +RUN --mount=type=cache,target=/root/.ccache \ + cmake --preset 'JetPack 6' \ + && cmake --build --parallel --preset 'JetPack 6' \ + && cmake --install build --component CUDA --strip --parallel 8 + +FROM base AS vulkan +RUN --mount=type=cache,target=/root/.ccache \ + cmake --preset 'Vulkan' \ + && cmake --build --parallel --preset 'Vulkan' \ + && cmake --install build --component Vulkan --strip --parallel 8 -RUN \ - wget -qO - https://packages.lunarg.com/lunarg-signing-key-pub.asc | gpg --dearmor -o /etc/apt/trusted.gpg.d/lunarg-signing-key-pub.gpg && \ - wget -qO /etc/apt/sources.list.d/lunarg-vulkan-${UBUNTU_VERSION}.list https://packages.lunarg.com/vulkan/${VULKAN_VER_BASE}/lunarg-vulkan-${VULKAN_VER_BASE}-${UBUNTU_VERSION}.list && \ - apt update && apt install -y vulkan-sdk +FROM base AS build +WORKDIR /go/src/github.com/ollama/ollama +COPY go.mod go.sum . +RUN curl -fsSL https://golang.org/dl/go$(awk '/^go/ { print $2 }' go.mod).linux-$(case $(uname -m) in x86_64) echo amd64 ;; aarch64) echo arm64 ;; esac).tar.gz | tar xz -C /usr/local +ENV PATH=/usr/local/go/bin:$PATH +RUN go mod download +COPY . . +ARG GOFLAGS="'-ldflags=-w -s'" +ENV CGO_ENABLED=1 +RUN --mount=type=cache,target=/root/.cache/go-build \ + go build -trimpath -buildmode=pie -o /bin/ollama . -# Last testet ollama-vulkan commit: -# 2d443b3dd660a1fd2760d64538512df93648b4bb -COPY patches/ /tmp/patches/ -RUN \ - git clone https://github.com/pufferffish/ollama-vulkan.git "/tmp/ollama-vulkan-git" && \ - cd "/tmp/ollama-vulkan-git" && \ - git checkout 2d443b3dd660a1fd2760d64538512df93648b4bb -b ollama_vulkan_stable && \ - git config user.name "Builder" && git config user.email "builder@local" && \ - git remote add ollama_vanilla https://github.com/ollama/ollama.git && \ - git fetch ollama_vanilla --tags && git checkout v0.5.13 -b ollama_vanilla_stable && \ - git checkout ollama_vanilla_stable && git merge ollama_vulkan_stable --allow-unrelated-histories --no-edit && \ - for p in /tmp/patches/*.patch; do patch -p1 < $p; done -RUN \ - cd "/tmp/ollama-vulkan-git" && \ - make -f Makefile.sync clean sync +FROM --platform=linux/amd64 scratch AS amd64 +COPY --from=cuda-11 dist/lib/ollama/cuda_v11 /lib/ollama/cuda_v11 +COPY --from=cuda-12 dist/lib/ollama/cuda_v12 /lib/ollama/cuda_v12 +COPY --from=vulkan dist/lib/ollama/vulkan /lib/ollama/vulkan +FROM --platform=linux/arm64 scratch AS arm64 +COPY --from=cuda-11 dist/lib/ollama/cuda_v11 /lib/ollama/cuda_v11 +COPY --from=cuda-12 dist/lib/ollama/cuda_v12 /lib/ollama/cuda_v12 +COPY --from=jetpack-5 dist/lib/ollama/cuda_v11 lib/ollama/cuda_jetpack5 +COPY --from=jetpack-6 dist/lib/ollama/cuda_v12 lib/ollama/cuda_jetpack6 -FROM builder AS cpu-build -RUN \ - cd "/tmp/ollama-vulkan-git" && \ - cmake --preset CPU && cmake --build --parallel --preset CPU && \ - cmake --install build --component CPU --strip +FROM scratch AS rocm +COPY --from=rocm-6 dist/lib/ollama/rocm /lib/ollama/rocm -FROM builder AS vulkan-build -RUN \ - cd "/tmp/ollama-vulkan-git" && \ - cmake --preset Vulkan && \ - cmake --build --parallel --preset Vulkan && \ - cmake --install build --component Vulkan --strip - -FROM builder AS binary-build -RUN \ - cd "/tmp/ollama-vulkan-git" && \ - . scripts/env.sh && \ - mkdir -p dist/bin && \ - go build -trimpath -buildmode=pie -o dist/bin/ollama . - - -FROM --platform=linux/amd64 library/ubuntu:noble -RUN \ - apt-get update && apt -y dist-upgrade && \ - apt-get install -y ca-certificates libcap2 libvulkan1 && \ - apt-get clean && rm -rf /var/lib/apt/lists/* - -# Install ROCm -RUN \ - apt update && \ - apt install -y wget python3-setuptools python3-wheel && \ - wget https://repo.radeon.com/amdgpu-install/6.3.3/ubuntu/noble/amdgpu-install_6.3.60303-1_all.deb -O /tmp/amdgpu-install_6.3.60303-1_all.deb && \ - apt install -y /tmp/amdgpu-install_6.3.60303-1_all.deb && \ - apt update && apt install -y rocm && \ - apt-get clean && rm -rf /var/lib/apt/lists/* - - -COPY --from=cpu-build /tmp/ollama-vulkan-git/dist/lib/ollama/ /lib/ollama/ -COPY --from=vulkan-build /tmp/ollama-vulkan-git/dist/lib/ollama/vulkan/ /lib/ollama/vulkan/ -COPY --from=binary-build /tmp/ollama-vulkan-git/dist/bin/ /bin/ - -RUN find /lib/ollama && find /bin/ollama +FROM ${FLAVOR} AS archive +ARG VULKANVERSION +COPY --from=cpu dist/lib/ollama /lib/ollama +COPY --from=build /bin/ollama /bin/ollama +FROM ubuntu:24.04 +RUN apt-get update \ + && apt-get install -y ca-certificates libcap2 libvulkan1 \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* +COPY --from=archive /bin /usr/bin +COPY --from=archive /lib/ollama /usr/lib/ollama +ENV PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +ENV LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64 +ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility +ENV NVIDIA_VISIBLE_DEVICES=all +ENV OLLAMA_HOST=0.0.0.0:11434 EXPOSE 11434 -ENV OLLAMA_HOST 0.0.0.0 - ENTRYPOINT ["/bin/ollama"] -CMD ["serve"] +CMD ["serve"] \ No newline at end of file From 9cb4ad02e243ef56ed1cf2ce8f4fd2285433cdf7 Mon Sep 17 00:00:00 2001 From: Vadim Grinco Date: Tue, 11 Mar 2025 14:34:17 +0100 Subject: [PATCH 29/29] This is no longer needed Signed-off-by: Vadim Grinco --- patches/00-fix-vulkan-building.patch | 15297 ------------------------- 1 file changed, 15297 deletions(-) delete mode 100644 patches/00-fix-vulkan-building.patch diff --git a/patches/00-fix-vulkan-building.patch b/patches/00-fix-vulkan-building.patch deleted file mode 100644 index 52e498ee2..000000000 --- a/patches/00-fix-vulkan-building.patch +++ /dev/null @@ -1,15297 +0,0 @@ -From 7c5f98c4cbfaf472a0d05baa3cc61afdcaeee7de Mon Sep 17 00:00:00 2001 -From: dream -Date: Thu, 13 Feb 2025 18:58:59 +0800 -Subject: [PATCH 2/2] fix: fix vulkan building - -1. Add preset for vulkan. -2. Add backend ggml-vulkan. -3. Add some log info. ---- - CMakePresets.json | 13 +- - discover/gpu.go | 7 +- - .../ggml/ggml/src/ggml-vulkan/CMakeLists.txt | 92 + - .../ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp | 8745 +++++++++++++++++ - .../ggml-vulkan/vulkan-shaders/CMakeLists.txt | 9 + - .../src/ggml-vulkan/vulkan-shaders/acc.comp | 29 + - .../src/ggml-vulkan/vulkan-shaders/add.comp | 29 + - .../ggml-vulkan/vulkan-shaders/argsort.comp | 69 + - .../src/ggml-vulkan/vulkan-shaders/clamp.comp | 17 + - .../ggml-vulkan/vulkan-shaders/concat.comp | 41 + - .../vulkan-shaders/contig_copy.comp | 42 + - .../src/ggml-vulkan/vulkan-shaders/copy.comp | 20 + - .../src/ggml-vulkan/vulkan-shaders/cos.comp | 17 + - .../vulkan-shaders/dequant_f32.comp | 20 + - .../vulkan-shaders/dequant_funcs.comp | 118 + - .../vulkan-shaders/dequant_funcs_cm2.comp | 325 + - .../vulkan-shaders/dequant_head.comp | 13 + - .../vulkan-shaders/dequant_iq4_nl.comp | 32 + - .../vulkan-shaders/dequant_q2_k.comp | 34 + - .../vulkan-shaders/dequant_q3_k.comp | 42 + - .../vulkan-shaders/dequant_q4_0.comp | 30 + - .../vulkan-shaders/dequant_q4_1.comp | 32 + - .../vulkan-shaders/dequant_q4_k.comp | 68 + - .../vulkan-shaders/dequant_q5_0.comp | 34 + - .../vulkan-shaders/dequant_q5_1.comp | 35 + - .../vulkan-shaders/dequant_q5_k.comp | 70 + - .../vulkan-shaders/dequant_q6_k.comp | 33 + - .../vulkan-shaders/dequant_q8_0.comp | 31 + - .../vulkan-shaders/diag_mask_inf.comp | 34 + - .../src/ggml-vulkan/vulkan-shaders/div.comp | 27 + - .../vulkan-shaders/flash_attn_cm2.comp | 289 + - .../src/ggml-vulkan/vulkan-shaders/gelu.comp | 25 + - .../vulkan-shaders/gelu_quick.comp | 23 + - .../vulkan-shaders/generic_binary_head.comp | 64 + - .../vulkan-shaders/generic_head.comp | 9 + - .../vulkan-shaders/generic_unary_head.comp | 56 + - .../ggml-vulkan/vulkan-shaders/get_rows.comp | 28 + - .../vulkan-shaders/get_rows_quant.comp | 39 + - .../vulkan-shaders/group_norm.comp | 66 + - .../ggml-vulkan/vulkan-shaders/im2col.comp | 87 + - .../vulkan-shaders/leaky_relu.comp | 22 + - .../src/ggml-vulkan/vulkan-shaders/mul.comp | 27 + - .../mul_mat_split_k_reduce.comp | 48 + - .../vulkan-shaders/mul_mat_vec.comp | 152 + - .../vulkan-shaders/mul_mat_vec_base.comp | 118 + - .../vulkan-shaders/mul_mat_vec_nc.comp | 71 + - .../vulkan-shaders/mul_mat_vec_p021.comp | 73 + - .../vulkan-shaders/mul_mat_vec_q2_k.comp | 115 + - .../vulkan-shaders/mul_mat_vec_q3_k.comp | 103 + - .../vulkan-shaders/mul_mat_vec_q4_k.comp | 133 + - .../vulkan-shaders/mul_mat_vec_q5_k.comp | 162 + - .../vulkan-shaders/mul_mat_vec_q6_k.comp | 112 + - .../ggml-vulkan/vulkan-shaders/mul_mm.comp | 631 ++ - .../vulkan-shaders/mul_mm_cm2.comp | 328 + - .../src/ggml-vulkan/vulkan-shaders/norm.comp | 44 + - .../src/ggml-vulkan/vulkan-shaders/pad.comp | 28 + - .../ggml-vulkan/vulkan-shaders/pool2d.comp | 74 + - .../src/ggml-vulkan/vulkan-shaders/relu.comp | 21 + - .../ggml-vulkan/vulkan-shaders/repeat.comp | 26 + - .../ggml-vulkan/vulkan-shaders/rms_norm.comp | 42 + - .../ggml-vulkan/vulkan-shaders/rope_head.comp | 49 + - .../ggml-vulkan/vulkan-shaders/rope_neox.comp | 37 + - .../ggml-vulkan/vulkan-shaders/rope_norm.comp | 37 + - .../src/ggml-vulkan/vulkan-shaders/scale.comp | 24 + - .../src/ggml-vulkan/vulkan-shaders/silu.comp | 22 + - .../src/ggml-vulkan/vulkan-shaders/sin.comp | 17 + - .../ggml-vulkan/vulkan-shaders/soft_max.comp | 174 + - .../ggml-vulkan/vulkan-shaders/square.comp | 17 + - .../ggml-vulkan/vulkan-shaders/sum_rows.comp | 37 + - .../src/ggml-vulkan/vulkan-shaders/tanh.comp | 20 + - .../vulkan-shaders/test_coopmat2_support.comp | 7 + - .../vulkan-shaders/timestep_embedding.comp | 41 + - .../src/ggml-vulkan/vulkan-shaders/types.comp | 323 + - .../ggml-vulkan/vulkan-shaders/upscale.comp | 36 + - .../vulkan-shaders/vulkan-shaders-gen.cpp | 594 ++ - .../src/ggml-vulkan/vulkan-shaders/wkv6.comp | 87 + - 76 files changed, 14642 insertions(+), 4 deletions(-) - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp - create mode 100644 ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp - -diff --git a/CMakePresets.json b/CMakePresets.json -index 3ecb0a8f..a77f15ba 100644 ---- a/CMakePresets.json -+++ b/CMakePresets.json -@@ -58,7 +58,11 @@ - "cacheVariables": { - "AMDGPU_TARGETS": "gfx803;gfx900;gfx940;gfx941;gfx942;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102;gfx906:xnack-;gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-" - } -- } -+ }, -+ { -+ "name": "Vulkan", -+ "inherits": [ "Default" ] -+ } - ], - "buildPresets": [ - { -@@ -105,6 +109,11 @@ - "name": "ROCm 6", - "inherits": [ "ROCm" ], - "configurePreset": "ROCm 6" -- } -+ }, -+ { -+ "name": "Vulkan", -+ "targets": [ "ggml-vulkan" ], -+ "configurePreset": "Vulkan" -+ } - ] - } -diff --git a/discover/gpu.go b/discover/gpu.go -index ec96f5d4..8079be99 100644 ---- a/discover/gpu.go -+++ b/discover/gpu.go -@@ -197,7 +197,10 @@ func initVulkanHandles() *vulkanHandles { - libcapPaths := FindLibCapLibs() - - if len(vulkanPaths) > 0 && len(libcapPaths) > 0 { -+ slog.Info("vulkan: load libvulkan and libcap ok") - vHandles.deviceCount, vHandles.vulkan, vulkanLibPath, libcapLibPath = LoadVulkanMgmt(vulkanPaths, libcapPaths) -+ } else { -+ slog.Info("vulkan: failed to load libvulkan or libcap") - } - - return vHandles -@@ -426,7 +429,7 @@ func GetGPUInfo() GpuInfoList { - gpuInfo.ID = C.GoString(&memInfo.gpu_id[0]) - gpuInfo.Compute = fmt.Sprintf("%d.%d", memInfo.major, memInfo.minor) - gpuInfo.MinimumMemory = 0 -- gpuInfo.DependencyPath = depPaths -+ gpuInfo.DependencyPath = []string{LibOllamaPath} - gpuInfo.Name = C.GoString(&memInfo.gpu_name[0]) - gpuInfo.DriverMajor = int(memInfo.major) - gpuInfo.DriverMinor = int(memInfo.minor) -@@ -768,7 +771,7 @@ func LoadVulkanMgmt(vulkanLibPaths []string, capLibPaths []string) (int, *C.vk_h - - C.vk_init(vkLib, capLib, &resp) - if resp.err != nil { -- slog.Debug("Unable to load vulkan", "library", vkLibPath, capLibPath, "error", C.GoString(resp.err)) -+ slog.Error("Unable to load vulkan", "library", vkLibPath, capLibPath, "error", C.GoString(resp.err)) - C.free(unsafe.Pointer(resp.err)) - } else { - return int(resp.num_devices), &resp.ch, vkLibPath, capLibPath -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt b/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt -new file mode 100644 -index 00000000..9501de73 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/CMakeLists.txt -@@ -0,0 +1,92 @@ -+find_package(Vulkan COMPONENTS glslc REQUIRED) -+ -+if (Vulkan_FOUND) -+ message(STATUS "Vulkan found") -+ -+ ggml_add_backend_library(ggml-vulkan -+ ggml-vulkan.cpp -+ ../../include/ggml-vulkan.h -+ ) -+ -+ # Compile a test shader to determine whether GL_NV_cooperative_matrix2 is supported. -+ # If it's not, there will be an error to stderr. -+ # If it's supported, set a define to indicate that we should compile those shaders -+ execute_process(COMMAND ${Vulkan_GLSLC_EXECUTABLE} -o - -fshader-stage=compute --target-env=vulkan1.3 "${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders/test_coopmat2_support.comp" -+ OUTPUT_VARIABLE glslc_output -+ ERROR_VARIABLE glslc_error) -+ -+ if (${glslc_error} MATCHES ".*extension not supported: GL_NV_cooperative_matrix2.*") -+ message(STATUS "GL_NV_cooperative_matrix2 not supported by glslc") -+ else() -+ message(STATUS "GL_NV_cooperative_matrix2 supported by glslc") -+ add_compile_definitions(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) -+ endif() -+ -+ target_link_libraries(ggml-vulkan PRIVATE Vulkan::Vulkan) -+ target_include_directories(ggml-vulkan PRIVATE ${CMAKE_CURRENT_BINARY_DIR}) -+ -+ # Workaround to the "can't dereference invalidated vector iterator" bug in clang-cl debug build -+ # Posssibly relevant: https://stackoverflow.com/questions/74748276/visual-studio-no-displays-the-correct-length-of-stdvector -+ if (MSVC AND CMAKE_CXX_COMPILER_ID STREQUAL "Clang") -+ add_compile_definitions(_ITERATOR_DEBUG_LEVEL=0) -+ endif() -+ -+ if (GGML_VULKAN_CHECK_RESULTS) -+ add_compile_definitions(GGML_VULKAN_CHECK_RESULTS) -+ endif() -+ -+ if (GGML_VULKAN_DEBUG) -+ add_compile_definitions(GGML_VULKAN_DEBUG) -+ endif() -+ -+ if (GGML_VULKAN_MEMORY_DEBUG) -+ add_compile_definitions(GGML_VULKAN_MEMORY_DEBUG) -+ endif() -+ -+ if (GGML_VULKAN_SHADER_DEBUG_INFO) -+ add_compile_definitions(GGML_VULKAN_SHADER_DEBUG_INFO) -+ endif() -+ -+ if (GGML_VULKAN_PERF) -+ add_compile_definitions(GGML_VULKAN_PERF) -+ endif() -+ -+ if (GGML_VULKAN_VALIDATE) -+ add_compile_definitions(GGML_VULKAN_VALIDATE) -+ endif() -+ -+ if (GGML_VULKAN_RUN_TESTS) -+ add_compile_definitions(GGML_VULKAN_RUN_TESTS) -+ endif() -+ -+ add_subdirectory(vulkan-shaders) -+ -+ set (_ggml_vk_genshaders_cmd vulkan-shaders-gen) -+ set (_ggml_vk_header ${CMAKE_CURRENT_BINARY_DIR}/ggml-vulkan-shaders.hpp) -+ set (_ggml_vk_source ${CMAKE_CURRENT_BINARY_DIR}/ggml-vulkan-shaders.cpp) -+ set (_ggml_vk_input_dir ${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders) -+ set (_ggml_vk_output_dir ${CMAKE_CURRENT_BINARY_DIR}/vulkan-shaders.spv) -+ -+ file(GLOB _ggml_vk_shader_deps "${_ggml_vk_input_dir}/*.comp") -+ -+ add_custom_command( -+ OUTPUT ${_ggml_vk_header} -+ ${_ggml_vk_source} -+ -+ COMMAND "$/${_ggml_vk_genshaders_cmd}" -+ --glslc ${Vulkan_GLSLC_EXECUTABLE} -+ --input-dir ${_ggml_vk_input_dir} -+ --output-dir ${_ggml_vk_output_dir} -+ --target-hpp ${_ggml_vk_header} -+ --target-cpp ${_ggml_vk_source} -+ --no-clean -+ -+ DEPENDS ${_ggml_vk_shader_deps} ${_ggml_vk_genshaders_cmd} -+ COMMENT "Generate vulkan shaders" -+ ) -+ -+ target_sources(ggml-vulkan PRIVATE ${_ggml_vk_source} ${_ggml_vk_header}) -+ -+else() -+ message(WARNING "Vulkan not found") -+endif() -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp -new file mode 100644 -index 00000000..d75cd6d6 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/ggml-vulkan.cpp -@@ -0,0 +1,8745 @@ -+#include "ggml-vulkan.h" -+#include -+#if defined(GGML_VULKAN_RUN_TESTS) || defined(GGML_VULKAN_PERF) || defined(GGML_VULKAN_CHECK_RESULTS) -+#include -+#include "ggml-cpu.h" -+#endif -+ -+#include -+ -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+ -+#include "ggml-impl.h" -+#include "ggml-backend-impl.h" -+ -+#include "ggml-vulkan-shaders.hpp" -+ -+#define VK_API_VERSION VK_API_VERSION_1_2 -+ -+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) -+ -+#define VK_VENDOR_ID_AMD 0x1002 -+#define VK_VENDOR_ID_APPLE 0x106b -+#define VK_VENDOR_ID_INTEL 0x8086 -+#define VK_VENDOR_ID_NVIDIA 0x10de -+ -+#define VK_DEVICE_DESCRIPTOR_POOL_SIZE 32 -+ -+#define GGML_VK_MAX_NODES 8192 -+ -+#define MAX_VK_BUFFERS 256 -+ -+#define VK_CHECK(err, msg) \ -+ do { \ -+ vk::Result err_ = (err); \ -+ if (err_ != vk::Result::eSuccess) { \ -+ fprintf(stderr, "ggml_vulkan: %s error %s at %s:%d\n", \ -+ #err, to_string(err_).c_str(), __FILE__, __LINE__); \ -+ exit(1); \ -+ } \ -+ } while (0) -+ -+#ifdef GGML_VULKAN_DEBUG -+#define VK_LOG_DEBUG(msg) std::cerr << msg << std::endl -+#else -+#define VK_LOG_DEBUG(msg) ((void) 0) -+#endif // GGML_VULKAN_DEBUG -+ -+struct ggml_backend_vk_context; -+ -+struct vk_queue { -+ uint32_t queue_family_index; -+ vk::Queue queue; -+ vk::CommandPool pool; -+ uint32_t cmd_buffer_idx; -+ std::vector cmd_buffers; -+ -+ vk::PipelineStageFlags stage_flags; -+ -+ bool transfer_only; -+}; -+ -+struct vk_pipeline_struct { -+ std::string name; -+ vk::ShaderModule shader_module; -+ vk::DescriptorSetLayout dsl; -+ std::vector descriptor_pools; -+ std::vector descriptor_sets; -+ uint32_t descriptor_set_idx; -+ vk::PipelineLayout layout; -+ vk::Pipeline pipeline; -+ uint32_t push_constant_size; -+ uint32_t parameter_count; -+ std::array wg_denoms; -+ uint32_t align; -+}; -+ -+typedef std::shared_ptr vk_pipeline; -+typedef std::weak_ptr vk_pipeline_ref; -+ -+static void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline); -+ -+struct vk_matmul_pipeline_struct { -+ vk_pipeline l, m, s; -+ vk_pipeline a_l, a_m, a_s; -+}; -+ -+typedef std::shared_ptr vk_matmul_pipeline; -+ -+struct vk_matmul_pipeline2 { -+ vk_matmul_pipeline2() { -+ f16acc = std::make_shared(); -+ f32acc = std::make_shared(); -+ } -+ vk_matmul_pipeline f32acc; -+ vk_matmul_pipeline f16acc; -+}; -+ -+struct vk_device_struct; -+typedef std::shared_ptr vk_device; -+typedef std::weak_ptr vk_device_ref; -+ -+struct vk_buffer_struct; -+typedef std::shared_ptr vk_buffer; -+typedef std::weak_ptr vk_buffer_ref; -+ -+struct ggml_backend_vk_buffer_type_context { -+ std::string name; -+ vk_device device; -+}; -+ -+static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft); -+static ggml_backend_buffer_t ggml_backend_vk_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size); -+static size_t ggml_backend_vk_buffer_type_get_alignment(ggml_backend_buffer_type_t buft); -+static size_t ggml_backend_vk_buffer_type_get_max_size(ggml_backend_buffer_type_t buft); -+static size_t ggml_backend_vk_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor); -+static ggml_backend_buffer_type_i ggml_backend_vk_buffer_type_interface = { -+ /* .get_name = */ ggml_backend_vk_buffer_type_name, -+ /* .alloc_buffer = */ ggml_backend_vk_buffer_type_alloc_buffer, -+ /* .get_alignment = */ ggml_backend_vk_buffer_type_get_alignment, -+ /* .get_max_size = */ ggml_backend_vk_buffer_type_get_max_size, -+ /* .get_alloc_size = */ ggml_backend_vk_buffer_type_get_alloc_size, -+ /* .is_host = */ NULL, -+}; -+ -+#ifdef GGML_VULKAN_MEMORY_DEBUG -+class vk_memory_logger; -+#endif -+#ifdef GGML_VULKAN_PERF -+class vk_perf_logger; -+#endif -+static void ggml_vk_destroy_buffer(vk_buffer& buf); -+ -+static constexpr uint32_t mul_mat_vec_max_cols = 8; -+ -+struct vk_device_struct { -+ std::mutex mutex; -+ -+ vk::PhysicalDevice physical_device; -+ vk::PhysicalDeviceProperties properties; -+ std::string name; -+ uint64_t max_memory_allocation_size; -+ bool fp16; -+ bool pipeline_robustness; -+ vk::Device device; -+ uint32_t vendor_id; -+ vk_queue compute_queue; -+ vk_queue transfer_queue; -+ bool single_queue; -+ uint32_t subgroup_size; -+ uint32_t shader_core_count; -+ bool uma; -+ bool float_controls_rte_fp16; -+ -+ bool subgroup_size_control; -+ uint32_t subgroup_min_size; -+ uint32_t subgroup_max_size; -+ bool subgroup_require_full_support; -+ -+ bool coopmat_support; -+ bool coopmat_acc_f32_support; -+ bool coopmat_acc_f16_support; -+ uint32_t coopmat_m; -+ uint32_t coopmat_n; -+ uint32_t coopmat_k; -+ bool coopmat2; -+ -+ size_t idx; -+ -+ bool mul_mat_l; -+ bool mul_mat_m; -+ bool mul_mat_s; -+ bool mul_mat_id_l; -+ bool mul_mat_id_m; -+ bool mul_mat_id_s; -+ -+ vk_matmul_pipeline pipeline_matmul_f32; -+ vk_matmul_pipeline pipeline_matmul_f32_f16; -+ vk_matmul_pipeline2 pipeline_matmul_f16; -+ vk_matmul_pipeline2 pipeline_matmul_f16_f32; -+ vk_pipeline pipeline_matmul_split_k_reduce; -+ -+ vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_COUNT]; -+ vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat[GGML_TYPE_COUNT]; -+ -+ vk_matmul_pipeline pipeline_matmul_id_f32; -+ vk_matmul_pipeline2 pipeline_matmul_id_f16; -+ vk_matmul_pipeline2 pipeline_matmul_id_f16_f32; -+ -+ vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id[GGML_TYPE_COUNT]; -+ -+ vk_pipeline pipeline_dequant[GGML_TYPE_COUNT]; -+ vk_pipeline pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_COUNT][mul_mat_vec_max_cols]; -+ vk_pipeline pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_COUNT][mul_mat_vec_max_cols]; -+ vk_pipeline pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_COUNT]; -+ -+ vk_pipeline pipeline_mul_mat_vec_p021_f16_f32; -+ vk_pipeline pipeline_mul_mat_vec_nc_f16_f32; -+ vk_pipeline pipeline_get_rows[GGML_TYPE_COUNT]; -+ vk_pipeline pipeline_get_rows_f32[GGML_TYPE_COUNT]; -+ vk_pipeline pipeline_acc_f32; -+ vk_pipeline pipeline_add_f32, pipeline_add_f32_norepeat; -+ vk_pipeline pipeline_add_f16_f32_f16, pipeline_add_f16_f32_f16_norepeat; -+ vk_pipeline pipeline_mul_f32, pipeline_mul_f32_norepeat; -+ vk_pipeline pipeline_div_f32, pipeline_div_f32_norepeat; -+ vk_pipeline pipeline_concat_f32, pipeline_concat_f16, pipeline_concat_i32; -+ vk_pipeline pipeline_upscale_f32; -+ vk_pipeline pipeline_scale_f32; -+ vk_pipeline pipeline_sqr_f32; -+ vk_pipeline pipeline_sin_f32; -+ vk_pipeline pipeline_cos_f32; -+ vk_pipeline pipeline_clamp_f32; -+ vk_pipeline pipeline_pad_f32; -+ vk_pipeline pipeline_repeat_f32; -+ vk_pipeline pipeline_cpy_f32_f32, pipeline_cpy_f32_f16, pipeline_cpy_f16_f16; -+ vk_pipeline pipeline_contig_cpy_f32_f32, pipeline_contig_cpy_f32_f16, pipeline_contig_cpy_f16_f16; -+ vk_pipeline pipeline_norm_f32; -+ vk_pipeline pipeline_group_norm_f32; -+ vk_pipeline pipeline_rms_norm_f32; -+ vk_pipeline pipeline_gelu_f32; -+ vk_pipeline pipeline_gelu_quick_f32; -+ vk_pipeline pipeline_silu_f32; -+ vk_pipeline pipeline_relu_f32; -+ vk_pipeline pipeline_leaky_relu_f32; -+ vk_pipeline pipeline_tanh_f32; -+ vk_pipeline pipeline_diag_mask_inf_f32; -+ vk_pipeline pipeline_soft_max_f32, pipeline_soft_max_f32_f16; -+ vk_pipeline pipeline_soft_max_f32_wg512, pipeline_soft_max_f32_f16_wg512; -+ vk_pipeline pipeline_rope_norm_f32, pipeline_rope_norm_f16; -+ vk_pipeline pipeline_rope_neox_f32, pipeline_rope_neox_f16; -+ vk_pipeline pipeline_argsort_f32; -+ vk_pipeline pipeline_sum_rows_f32; -+ vk_pipeline pipeline_im2col_f32, pipeline_im2col_f32_f16; -+ vk_pipeline pipeline_timestep_embedding_f32; -+ vk_pipeline pipeline_pool2d_f32; -+ vk_pipeline pipeline_rwkv_wkv6_f32; -+ -+ // [2][2][2] is for {f16acc,f32acc}x{large,small_rows}x{unaligned, aligned} -+ vk_pipeline pipeline_flash_attn_f32_f16_D64[GGML_TYPE_COUNT][2][2][2]; -+ vk_pipeline pipeline_flash_attn_f32_f16_D80[GGML_TYPE_COUNT][2][2][2]; -+ vk_pipeline pipeline_flash_attn_f32_f16_D96[GGML_TYPE_COUNT][2][2][2]; -+ vk_pipeline pipeline_flash_attn_f32_f16_D112[GGML_TYPE_COUNT][2][2][2]; -+ vk_pipeline pipeline_flash_attn_f32_f16_D128[GGML_TYPE_COUNT][2][2][2]; -+ vk_pipeline pipeline_flash_attn_f32_f16_D256[GGML_TYPE_COUNT][2][2][2]; -+ -+ std::unordered_map pipelines; -+ std::unordered_map pipeline_descriptor_set_requirements; -+ -+ std::vector> pinned_memory; -+ -+ vk::Fence fence; -+ vk_buffer sync_staging; -+ -+ ggml_backend_buffer_type buffer_type; -+ -+#ifdef GGML_VULKAN_MEMORY_DEBUG -+ std::unique_ptr memory_logger; -+#endif -+#ifdef GGML_VULKAN_PERF -+ std::unique_ptr perf_logger; -+#endif -+ -+ ~vk_device_struct() { -+ VK_LOG_DEBUG("destroy device " << name); -+ -+ device.destroyFence(fence); -+ -+ ggml_vk_destroy_buffer(sync_staging); -+ -+ device.destroyCommandPool(compute_queue.pool); -+ if (!single_queue) { -+ device.destroyCommandPool(transfer_queue.pool); -+ } -+ -+ for (auto& pipeline : pipelines) { -+ if (pipeline.second.expired()) { -+ continue; -+ } -+ -+ vk_pipeline pl = pipeline.second.lock(); -+ ggml_vk_destroy_pipeline(device, pl); -+ } -+ pipelines.clear(); -+ -+ device.destroy(); -+ } -+}; -+ -+struct vk_buffer_struct { -+ vk::Buffer buffer = VK_NULL_HANDLE; -+ vk::DeviceMemory device_memory = VK_NULL_HANDLE; -+ vk::MemoryPropertyFlags memory_property_flags; -+ void * ptr; -+ size_t size = 0; -+ -+ vk_device device; -+ -+ ~vk_buffer_struct() { -+ if (size == 0) { -+ return; -+ } -+ VK_LOG_DEBUG("~vk_buffer_struct(" << buffer << ", " << size << ")"); -+ -+ device->device.freeMemory(device_memory); -+ device->device.destroyBuffer(buffer); -+ } -+}; -+ -+struct vk_subbuffer { -+ vk_buffer buffer; -+ uint64_t offset; -+ uint64_t size; -+ -+ operator vk::DescriptorBufferInfo() const { -+ return { buffer->buffer, offset, size }; -+ } -+}; -+ -+struct vk_semaphore { -+ vk::Semaphore s; -+ uint64_t value; -+}; -+ -+struct vk_submission { -+ vk::CommandBuffer buffer; -+ std::vector wait_semaphores; -+ std::vector signal_semaphores; -+}; -+ -+typedef std::vector vk_sequence; -+ -+struct vk_mat_mat_push_constants { -+ uint32_t M; uint32_t N; uint32_t K; -+ uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; -+ uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; -+ uint32_t k_split; -+ uint32_t ne02; uint32_t ne12; uint32_t broadcast2; uint32_t broadcast3; -+}; -+struct vk_mat_vec_push_constants { -+ uint32_t ncols; uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; -+ uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; -+ uint32_t ne02; uint32_t ne12; uint32_t broadcast2; uint32_t broadcast3; -+}; -+ -+struct vk_mat_mat_id_push_constants { -+ uint32_t M; uint32_t N; uint32_t K; -+ uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; -+ uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; -+ uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11; -+}; -+struct vk_mat_vec_id_push_constants { -+ uint32_t ncols; uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; -+ uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; -+ uint32_t nei0; uint32_t ne11; -+}; -+ -+struct vk_flash_attn_push_constants { -+ uint32_t N; -+ uint32_t KV; -+ -+ uint32_t ne1; -+ uint32_t ne2; -+ uint32_t ne3; -+ -+ uint32_t neq2; -+ uint32_t neq3; -+ uint32_t nek2; -+ uint32_t nek3; -+ uint32_t nev2; -+ uint32_t nev3; -+ uint32_t nem1; -+ -+ uint32_t nb02; -+ uint32_t nb03; -+ uint32_t nb12; -+ uint32_t nb13; -+ uint32_t nb22; -+ uint32_t nb23; -+ uint32_t nb31; -+ -+ float scale; -+ float max_bias; -+ float logit_softcap; -+ -+ uint32_t mask; -+ uint32_t n_head_log2; -+ float m0; -+ float m1; -+}; -+ -+struct vk_op_push_constants { -+ uint32_t KX; -+ uint32_t KY; -+ float param1; -+ float param2; -+}; -+ -+struct vk_op_unary_push_constants { -+ uint32_t ne; -+ uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; -+ uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; uint32_t nb10; uint32_t nb11; uint32_t nb12; uint32_t nb13; -+ uint32_t misalign_offsets; -+ float param1; float param2; -+ uint32_t ne0_012mp; uint32_t ne0_012L; -+ uint32_t ne0_01mp; uint32_t ne0_01L; -+ uint32_t ne0_0mp; uint32_t ne0_0L; -+ uint32_t ne1_012mp; uint32_t ne1_012L; -+ uint32_t ne1_01mp; uint32_t ne1_01L; -+ uint32_t ne1_0mp; uint32_t ne1_0L; -+}; -+static_assert(sizeof(vk_op_unary_push_constants) <= 128, "sizeof(vk_op_unary_push_constants) must be <= 128"); -+ -+// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1. -+// Precompute mp (m' in the paper) and L such that division -+// can be computed using a multiply (high 32b of 64b result) -+// and a shift: -+// -+// n/d = (mulhi(n, mp) + n) >> L; -+static void init_fastdiv_values(uint32_t d, uint32_t &mp, uint32_t &L) -+{ -+ // compute L = ceil(log2(d)); -+ L = 0; -+ while (L < 32 && (uint32_t{1} << L) < d) { -+ L++; -+ } -+ -+ mp = (uint32_t)((uint64_t{1} << 32) * ((uint64_t{1} << L) - d) / d + 1); -+} -+ -+template void init_pushconst_fastdiv(T &p) { -+ GGML_UNUSED(p); -+ static_assert(!std::is_const::value, "unexpected type"); -+} -+ -+template <> void init_pushconst_fastdiv(vk_op_unary_push_constants &p) { -+ // Compute magic values to divide by these six numbers. -+ init_fastdiv_values(p.ne02*p.ne01*p.ne00, p.ne0_012mp, p.ne0_012L); -+ init_fastdiv_values(p.ne01*p.ne00, p.ne0_01mp, p.ne0_01L); -+ init_fastdiv_values(p.ne00, p.ne0_0mp, p.ne0_0L); -+ init_fastdiv_values(p.ne12*p.ne11*p.ne10, p.ne1_012mp, p.ne1_012L); -+ init_fastdiv_values(p.ne11*p.ne10, p.ne1_01mp, p.ne1_01L); -+ init_fastdiv_values(p.ne10, p.ne1_0mp, p.ne1_0L); -+} -+ -+struct vk_op_binary_push_constants { -+ uint32_t ne; -+ uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; -+ uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; uint32_t nb10; uint32_t nb11; uint32_t nb12; uint32_t nb13; -+ uint32_t ne20; uint32_t ne21; uint32_t ne22; uint32_t ne23; uint32_t nb20; uint32_t nb21; uint32_t nb22; uint32_t nb23; -+ uint32_t misalign_offsets; -+ float param1; float param2; int32_t param3; -+}; -+ -+struct vk_op_diag_mask_push_constants { -+ uint32_t ncols; -+ uint32_t rows_per_channel; -+ int32_t n_past; -+}; -+ -+struct vk_op_rope_push_constants { -+ uint32_t ncols; -+ uint32_t n_dims; -+ float freq_scale; -+ uint32_t p_delta_rows; -+ float freq_base; -+ float ext_factor; -+ float attn_factor; -+ float corr_dims[2]; -+ float theta_scale; -+ uint32_t has_ff; -+}; -+ -+struct vk_op_soft_max_push_constants { -+ uint32_t KX; -+ uint32_t KY; -+ float scale; -+ float max_bias; -+ float m0; -+ float m1; -+ uint32_t n_head_log2; -+ uint32_t nrows_x; -+}; -+ -+struct vk_op_argsort_push_constants { -+ uint32_t ncols; -+ uint32_t ncols_pad; -+ int32_t order; -+}; -+ -+struct vk_op_im2col_push_constants { -+ uint32_t batch_offset; uint32_t offset_delta; -+ uint32_t IC; -+ uint32_t IW; uint32_t IH; -+ uint32_t OW; uint32_t OH; -+ uint32_t KW; uint32_t KH; -+ uint32_t pelements; -+ uint32_t CHW; -+ int32_t s0; int32_t s1; -+ int32_t p0; int32_t p1; -+ int32_t d0; int32_t d1; -+}; -+ -+struct vk_op_timestep_embedding_push_constants { -+ uint32_t nb1; -+ uint32_t dim; -+ uint32_t max_period; -+}; -+ -+struct vk_op_pool2d_push_constants { -+ uint32_t IW; uint32_t IH; -+ uint32_t OW; uint32_t OH; -+ uint32_t OC; -+ uint32_t pelements; -+ uint32_t op; -+ int32_t k0; int32_t k1; -+ int32_t s0; int32_t s1; -+ int32_t p0; int32_t p1; -+}; -+ -+struct vk_op_rwkv_wkv6_push_constants { -+ uint32_t B; -+ uint32_t T; -+ uint32_t C; -+ uint32_t H; -+}; -+ -+// Allow pre-recording command buffers -+struct vk_staging_memcpy { -+ vk_staging_memcpy(void * _dst, const void * _src, size_t _n) : dst(_dst), src(_src), n(_n) {} -+ -+ void * dst; -+ const void * src; -+ size_t n; -+}; -+ -+struct vk_op_upscale_push_constants { -+ uint32_t ne; uint32_t a_offset; uint32_t d_offset; -+ uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; -+ uint32_t ne10; uint32_t ne11; uint32_t ne12; uint32_t ne13; -+ float sf0; float sf1; float sf2; float sf3; -+}; -+ -+struct vk_context_struct { -+ vk_submission * s; -+ std::vector seqs; -+ -+ int exit_tensor_idx; -+ -+ std::vector in_memcpys; -+ std::vector out_memcpys; -+ -+ vk_queue * q; -+}; -+typedef std::shared_ptr vk_context; -+typedef std::weak_ptr vk_context_ref; -+ -+struct ggml_vk_garbage_collector { -+ std::vector tl_semaphores; -+ std::vector semaphores; -+ std::vector events; -+ std::vector temp_buffers; -+ std::vector contexts; -+}; -+ -+#if defined(GGML_VULKAN_MEMORY_DEBUG) || defined(GGML_VULKAN_DEBUG) -+#define VK_LOG_MEMORY(msg) std::cerr << "ggml_vulkan memory: " << msg << std::endl -+ -+static std::string format_size(size_t size) { -+ const size_t kib = 1024; -+ const size_t mib = kib * 1024; -+ const size_t gib = mib * 1024; -+ -+ std::ostringstream oss; -+ oss << std::fixed << std::setprecision(2); -+ -+ if (size >= gib) { -+ oss << static_cast(size) / gib << " GiB"; -+ } else if (size >= mib) { -+ oss << static_cast(size) / mib << " MiB"; -+ } else if (size >= kib) { -+ oss << static_cast(size) / kib << " KiB"; -+ } else { -+ oss << size << " B"; -+ } -+ -+ return oss.str(); -+} -+ -+static std::mutex log_mutex; -+ -+class vk_memory_logger { -+public: -+ vk_memory_logger(): total_device(0), total_host(0) {} -+ void log_allocation(vk_buffer_ref buf_ref, size_t size); -+ void log_deallocation(vk_buffer_ref buf_ref); -+ -+private: -+ std::map allocations; // Track allocations -+ size_t total_device; -+ size_t total_host; -+}; -+#else -+#define VK_LOG_MEMORY(msg) ((void) 0) -+#endif // GGML_VULKAN_MEMORY_DEBUG -+ -+#if defined(GGML_VULKAN_PERF) -+ -+class vk_perf_logger { -+public: -+ void print_timings() { -+ std::cerr << "----------------\nVulkan Timings:" << std::endl; -+ for (const auto& t : timings) { -+ uint64_t total = 0; -+ for (const auto& time : t.second) { -+ total += time; -+ } -+ std::cerr << t.first << ": " << t.second.size() << " x " << (total / t.second.size() / 1000.0) << " ms" << std::endl; -+ } -+ -+ timings.clear(); -+ } -+ -+ void log_timing(const ggml_tensor * node, uint64_t time) { -+ if (node->op == GGML_OP_UNARY) { -+ timings[ggml_unary_op_name(ggml_get_unary_op(node))].push_back(time); -+ return; -+ } -+ if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { -+ const uint64_t m = node->src[0]->ne[1]; -+ const uint64_t n = node->src[1]->ne[1]; -+ const uint64_t k = node->src[1]->ne[0]; -+ std::string name = ggml_op_name(node->op); -+ if (n == 1) { -+ name += "_VEC m=" + std::to_string(m) + " k=" + std::to_string(k); -+ } else { -+ name += " m=" + std::to_string(m) + " n=" + std::to_string(n) + " k=" + std::to_string(k); -+ } -+ timings[name].push_back(time); -+ return; -+ } -+ timings[ggml_op_name(node->op)].push_back(time); -+ } -+private: -+ std::map> timings; -+}; -+#endif // GGML_VULKAN_PERF -+ -+struct ggml_backend_vk_context { -+ std::string name; -+ -+ vk_device device; -+ -+ size_t semaphore_idx, event_idx; -+ ggml_vk_garbage_collector gc; -+ size_t prealloc_size_x, prealloc_size_y, prealloc_size_split_k; -+ vk_buffer prealloc_x, prealloc_y, prealloc_split_k; -+ vk::Fence fence; -+ -+ vk_buffer buffer_pool[MAX_VK_BUFFERS]; -+ -+ vk_context_ref compute_ctx; -+ vk_context_ref transfer_ctx; -+ -+ std::vector tensor_ctxs; -+}; -+ -+static void * const vk_ptr_base = (void *)(uintptr_t) 0x1000; // NOLINT -+ -+static uint64_t vk_tensor_offset(const ggml_tensor * tensor) { -+ if (tensor->view_src) { -+ return (uint8_t *) tensor->view_src->data - (uint8_t *) vk_ptr_base; -+ } -+ return (uint8_t *) tensor->data - (uint8_t *) vk_ptr_base; -+} -+ -+struct ggml_backend_vk_buffer_context { -+ vk_device_ref device; -+ vk_buffer dev_buffer; -+ std::string name; -+ -+ ggml_backend_vk_buffer_context(vk_device_ref device, vk_buffer&& dev_buffer, std::string& name) : -+ device(device), -+ dev_buffer(dev_buffer), -+ name(name) { -+ } -+ -+ ~ggml_backend_vk_buffer_context() { -+ ggml_vk_destroy_buffer(dev_buffer); -+ } -+}; -+ -+#ifdef GGML_VULKAN_MEMORY_DEBUG -+void vk_memory_logger::log_allocation(vk_buffer_ref buf_ref, size_t size) { -+ std::lock_guard guard(log_mutex); -+ vk_buffer buf = buf_ref.lock(); -+ const bool device = bool(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eDeviceLocal); -+ const std::string type = device ? "device" : "host"; -+ allocations[buf->buffer] = size; -+ total_device += device ? size : 0; -+ total_host += device ? 0 : size; -+ VK_LOG_MEMORY(buf->device->name << ": +" << format_size(size) << " " << type << " at " << buf->buffer << ". Total device: " << format_size(total_device) << ", total host: " << format_size(total_host)); -+} -+ -+void vk_memory_logger::log_deallocation(vk_buffer_ref buf_ref) { -+ if (buf_ref.expired() || buf_ref.lock()->size == 0) { -+ return; -+ } -+ -+ std::lock_guard guard(log_mutex); -+ vk_buffer buf = buf_ref.lock(); -+ const bool device = bool(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eDeviceLocal); -+ std::string type = device ? "device" : "host"; -+ auto it = allocations.find(buf->buffer); -+ total_device -= device ? it->second : 0; -+ total_host -= device ? 0 : it->second; -+ if (it != allocations.end()) { -+ VK_LOG_MEMORY(buf->device->name << ": -" << format_size(it->second) << " " << type << " at " << buf->buffer << ". Total device: " << format_size(total_device) << ", total host: " << format_size(total_host)); -+ allocations.erase(it); -+ } else { -+ VK_LOG_MEMORY("ERROR " << buf->device->name << ": Attempted to deallocate unknown " << type << " memory at " << buf->buffer); -+ } -+} -+#endif // GGML_VULKAN_MEMORY_DEBUG -+ -+struct vk_instance_t { -+ vk::Instance instance; -+ -+ std::vector device_indices; -+ vk_device devices[GGML_VK_MAX_DEVICES]; -+}; -+ -+static bool vk_instance_initialized = false; -+static vk_instance_t vk_instance; -+ -+#ifdef GGML_VULKAN_CHECK_RESULTS -+static size_t vk_skip_checks; -+static size_t vk_output_tensor; -+ -+static void ggml_vk_print_tensor(const ggml_tensor * tensor, const char * name); -+static void ggml_vk_check_results_0(ggml_tensor * tensor); -+static void ggml_vk_check_results_1(ggml_tensor * tensor); -+#endif -+ -+typedef void (*ggml_vk_func_t)(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); -+ -+static void ggml_backend_vk_free(ggml_backend_t backend); -+ -+// variables to track number of compiles in progress -+static uint32_t compile_count = 0; -+static std::mutex compile_count_mutex; -+static std::condition_variable compile_count_cond; -+ -+static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipeline, const std::string name, size_t spv_size, const void* spv_data, const std::string entrypoint, -+ uint32_t parameter_count, uint32_t push_constant_size, std::array wg_denoms, std::vector specialization_constants, -+ uint32_t align, bool disable_robustness, bool require_full_subgroups, uint32_t required_subgroup_size) { -+ VK_LOG_DEBUG("ggml_vk_create_pipeline(" << device->name << ", " << name << ", " << entrypoint << ", " << parameter_count << ", " << push_constant_size << -+ ", (" << wg_denoms[0] << "," << wg_denoms[1] << "," << wg_denoms[2] << "), specialization_constants, " << align << -+ ", " << disable_robustness << ", " << require_full_subgroups << ", " << required_subgroup_size << ")"); -+ GGML_ASSERT(parameter_count > 0); -+ GGML_ASSERT(wg_denoms[0] > 0 && wg_denoms[1] > 0 && wg_denoms[2] > 0); // NOLINT -+ -+ pipeline = std::make_shared(); -+ pipeline->name = name; -+ pipeline->parameter_count = parameter_count; -+ pipeline->push_constant_size = push_constant_size; -+ pipeline->wg_denoms = wg_denoms; -+ pipeline->align = align; -+ -+ vk::ShaderModuleCreateInfo shader_module_create_info({}, spv_size, reinterpret_cast(spv_data)); -+ pipeline->shader_module = device->device.createShaderModule(shader_module_create_info); -+ -+ std::vector dsl_binding; -+ std::vector dsl_binding_flags; -+ for (uint32_t i = 0; i < parameter_count; i++) { -+ dsl_binding.push_back({i, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute}); -+ dsl_binding_flags.push_back({}); -+ } -+ -+ vk::DescriptorSetLayoutBindingFlagsCreateInfo dslbfci = { dsl_binding_flags }; -+ -+ vk::PushConstantRange pcr( -+ vk::ShaderStageFlagBits::eCompute, -+ 0, -+ pipeline->push_constant_size -+ ); -+ -+ vk::DescriptorSetLayoutCreateInfo descriptor_set_layout_create_info( -+ {}, -+ dsl_binding); -+ descriptor_set_layout_create_info.setPNext(&dslbfci); -+ pipeline->dsl = device->device.createDescriptorSetLayout(descriptor_set_layout_create_info); -+ -+ vk::DescriptorPoolSize descriptor_pool_size(vk::DescriptorType::eStorageBuffer, pipeline->parameter_count * VK_DEVICE_DESCRIPTOR_POOL_SIZE); -+ vk::DescriptorPoolCreateInfo descriptor_pool_create_info({}, VK_DEVICE_DESCRIPTOR_POOL_SIZE, descriptor_pool_size); -+ pipeline->descriptor_pools.push_back(device->device.createDescriptorPool(descriptor_pool_create_info)); -+ -+ pipeline->descriptor_set_idx = 0; -+ -+ vk::PipelineLayoutCreateInfo pipeline_layout_create_info(vk::PipelineLayoutCreateFlags(), pipeline->dsl, pcr); -+ pipeline->layout = device->device.createPipelineLayout(pipeline_layout_create_info); -+ -+ std::vector specialization_entries(specialization_constants.size()); -+ -+ for (size_t i = 0; i < specialization_constants.size(); i++) { -+ specialization_entries[i].constantID = i; -+ specialization_entries[i].offset = i * sizeof(uint32_t); -+ specialization_entries[i].size = sizeof(uint32_t); -+ } -+ -+ vk::SpecializationInfo specialization_info( -+ specialization_entries.size(), -+ specialization_entries.data(), -+ specialization_constants.size() * sizeof(uint32_t), -+ specialization_constants.data() -+ ); -+ -+ vk::PipelineShaderStageCreateFlags pipeline_shader_stage_create_flags{}; -+ -+ if (device->subgroup_require_full_support && require_full_subgroups) { -+ pipeline_shader_stage_create_flags |= vk::PipelineShaderStageCreateFlagBits::eRequireFullSubgroupsEXT; -+ } -+ -+ vk::PipelineShaderStageCreateInfo pipeline_shader_create_info( -+ pipeline_shader_stage_create_flags, -+ vk::ShaderStageFlagBits::eCompute, -+ pipeline->shader_module, -+ entrypoint.c_str(), -+ &specialization_info); -+ -+ vk::PipelineShaderStageRequiredSubgroupSizeCreateInfoEXT pipeline_shader_stage_required_subgroup_size_create_info; -+ pipeline_shader_stage_required_subgroup_size_create_info.requiredSubgroupSize = required_subgroup_size; -+ if (device->subgroup_size_control && required_subgroup_size > 0) { -+ GGML_ASSERT(device->subgroup_min_size <= required_subgroup_size && required_subgroup_size <= device->subgroup_max_size); -+ pipeline_shader_create_info.setPNext(&pipeline_shader_stage_required_subgroup_size_create_info); -+ } -+ -+ vk::ComputePipelineCreateInfo compute_pipeline_create_info( -+ vk::PipelineCreateFlags{}, -+ pipeline_shader_create_info, -+ pipeline->layout); -+ -+ vk::PipelineRobustnessCreateInfoEXT rci; -+ -+ if (device->pipeline_robustness && disable_robustness) { -+ rci.storageBuffers = vk::PipelineRobustnessBufferBehaviorEXT::eDisabled; -+ rci.uniformBuffers = vk::PipelineRobustnessBufferBehaviorEXT::eDisabled; -+ compute_pipeline_create_info.setPNext(&rci); -+ } -+ -+ pipeline->pipeline = device->device.createComputePipeline(VK_NULL_HANDLE, compute_pipeline_create_info).value; -+ -+ { -+ std::lock_guard guard(device->mutex); -+ device->pipelines.insert({ pipeline->name, pipeline }); -+ } -+ -+ { -+ std::lock_guard guard(compile_count_mutex); -+ assert(compile_count > 0); -+ compile_count--; -+ -+ // "Progress bar" for shader compiles -+ static uint32_t total_compile_count = 0; -+ if ((total_compile_count++ % 10) == 0) { -+ std::cerr << "."; -+ } -+ } -+ compile_count_cond.notify_all(); -+} -+ -+static void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline) { -+ VK_LOG_DEBUG("ggml_pipeline_destroy_pipeline(" << pipeline->name << ")"); -+ for (auto& pool : pipeline->descriptor_pools) { -+ device.destroyDescriptorPool(pool); -+ } -+ pipeline->descriptor_pools.clear(); -+ pipeline->descriptor_sets.clear(); -+ pipeline->descriptor_set_idx = 0; -+ -+ device.destroyDescriptorSetLayout(pipeline->dsl); -+ -+ device.destroyPipelineLayout(pipeline->layout); -+ -+ device.destroyShaderModule(pipeline->shader_module); -+ -+ device.destroyPipeline(pipeline->pipeline); -+} -+ -+static void ggml_pipeline_request_descriptor_sets(vk_device& device, vk_pipeline& pipeline, uint32_t n) { -+ VK_LOG_DEBUG("ggml_pipeline_request_descriptor_sets(" << pipeline->name << ", " << n << ")"); -+ device->pipeline_descriptor_set_requirements[pipeline->name] += n; -+} -+ -+static void ggml_pipeline_allocate_descriptor_sets(vk_device& device) { -+ std::lock_guard guard(device->mutex); -+ -+ for (auto& pair : device->pipeline_descriptor_set_requirements) { -+ vk_pipeline pipeline = device->pipelines.at(pair.first).lock(); -+ const uint64_t n = pair.second; -+ -+ VK_LOG_DEBUG("ggml_pipeline_allocate_descriptor_sets(" << pipeline->name << ", " << n << ")"); -+ -+ if (pipeline->descriptor_sets.size() >= pipeline->descriptor_set_idx + n) { -+ // Enough descriptors are available -+ continue; -+ } -+ -+ uint32_t to_alloc = pipeline->descriptor_set_idx + n - pipeline->descriptor_sets.size(); -+ uint32_t pool_remaining = VK_DEVICE_DESCRIPTOR_POOL_SIZE - pipeline->descriptor_sets.size() % VK_DEVICE_DESCRIPTOR_POOL_SIZE; -+ uint32_t pool_idx = pipeline->descriptor_sets.size() / VK_DEVICE_DESCRIPTOR_POOL_SIZE; -+ -+ while (to_alloc > 0) { -+ const uint32_t alloc_count = std::min(pool_remaining, to_alloc); -+ to_alloc -= alloc_count; -+ pool_remaining = VK_DEVICE_DESCRIPTOR_POOL_SIZE; -+ -+ if (pool_idx >= pipeline->descriptor_pools.size()) { -+ vk::DescriptorPoolSize descriptor_pool_size(vk::DescriptorType::eStorageBuffer, pipeline->parameter_count * VK_DEVICE_DESCRIPTOR_POOL_SIZE); -+ vk::DescriptorPoolCreateInfo descriptor_pool_create_info({}, VK_DEVICE_DESCRIPTOR_POOL_SIZE, descriptor_pool_size); -+ pipeline->descriptor_pools.push_back(device->device.createDescriptorPool(descriptor_pool_create_info)); -+ } -+ -+ std::vector layouts(alloc_count); -+ for (uint32_t i = 0; i < alloc_count; i++) { -+ layouts[i] = pipeline->dsl; -+ } -+ vk::DescriptorSetAllocateInfo descriptor_set_alloc_info(pipeline->descriptor_pools[pool_idx], alloc_count, layouts.data()); -+ std::vector sets = device->device.allocateDescriptorSets(descriptor_set_alloc_info); -+ pipeline->descriptor_sets.insert(pipeline->descriptor_sets.end(), sets.begin(), sets.end()); -+ -+ pool_idx++; -+ } -+ } -+} -+ -+static void ggml_pipeline_cleanup(vk_pipeline& pipeline) { -+ VK_LOG_DEBUG("ggml_pipeline_cleanup(" << pipeline->name << ")"); -+ pipeline->descriptor_set_idx = 0; -+} -+ -+static vk::CommandBuffer ggml_vk_create_cmd_buffer(vk_device& device, vk_queue& q) { -+ VK_LOG_DEBUG("ggml_vk_create_cmd_buffer()"); -+ std::lock_guard guard(device->mutex); -+ -+ if (q.cmd_buffers.size() > q.cmd_buffer_idx) { -+ // Reuse command buffer -+ return q.cmd_buffers[q.cmd_buffer_idx++]; -+ } -+ -+ vk::CommandBufferAllocateInfo command_buffer_alloc_info( -+ q.pool, -+ vk::CommandBufferLevel::ePrimary, -+ 1); -+ const std::vector cmd_buffers = device->device.allocateCommandBuffers(command_buffer_alloc_info); -+ auto buf = cmd_buffers.front(); -+ -+ q.cmd_buffers.push_back(buf); -+ q.cmd_buffer_idx++; -+ -+ return buf; -+} -+ -+static vk_submission ggml_vk_create_submission(vk_device& device, vk_queue& q, std::vector wait_semaphores, std::vector signal_semaphores) { -+ VK_LOG_DEBUG("ggml_vk_create_submission()"); -+ vk_submission s; -+ s.buffer = ggml_vk_create_cmd_buffer(device, q); -+ s.wait_semaphores = std::move(wait_semaphores); -+ s.signal_semaphores = std::move(signal_semaphores); -+ return s; -+} -+ -+static void ggml_vk_submit(vk_context& ctx, vk::Fence fence) { -+ if (ctx->seqs.empty()) { -+ if (fence) { -+ ctx->q->queue.submit({}, fence); -+ } -+ return; -+ } -+ VK_LOG_DEBUG("ggml_vk_submit(" << ctx << ", " << fence << ")"); -+ -+ std::vector> tl_wait_vals; -+ std::vector> tl_signal_vals; -+ std::vector> tl_wait_semaphores; -+ std::vector> tl_signal_semaphores; -+ std::vector tl_submit_infos; -+ std::vector submit_infos; -+ int idx = -1; -+ std::vector> stage_flags; -+ -+ size_t reserve = 0; -+ -+ for (const auto& sequence : ctx->seqs) { -+ reserve += sequence.size(); -+ } -+ -+ // Pre-reserve vectors to prevent reallocation, which invalidates pointers -+ tl_wait_semaphores.reserve(reserve); -+ tl_wait_vals.reserve(reserve); -+ tl_signal_semaphores.reserve(reserve); -+ tl_signal_vals.reserve(reserve); -+ tl_submit_infos.reserve(reserve); -+ submit_infos.reserve(reserve); -+ stage_flags.reserve(reserve); -+ -+ for (const auto& sequence : ctx->seqs) { -+ for (const auto& submission : sequence) { -+ stage_flags.push_back({}); -+ idx++; -+ tl_wait_vals.push_back({}); -+ tl_wait_semaphores.push_back({}); -+ tl_signal_vals.push_back({}); -+ tl_signal_semaphores.push_back({}); -+ for (size_t i = 0; i < submission.wait_semaphores.size(); i++) { -+ stage_flags[idx].push_back(ctx->q->stage_flags); -+ tl_wait_vals[idx].push_back(submission.wait_semaphores[i].value); -+ tl_wait_semaphores[idx].push_back(submission.wait_semaphores[i].s); -+ } -+ for (size_t i = 0; i < submission.signal_semaphores.size(); i++) { -+ tl_signal_vals[idx].push_back(submission.signal_semaphores[i].value); -+ tl_signal_semaphores[idx].push_back(submission.signal_semaphores[i].s); -+ } -+ tl_submit_infos.push_back({ -+ (uint32_t) submission.wait_semaphores.size(), -+ tl_wait_vals[idx].data(), -+ (uint32_t) submission.signal_semaphores.size(), -+ tl_signal_vals[idx].data(), -+ }); -+ tl_submit_infos[idx].sType = vk::StructureType::eTimelineSemaphoreSubmitInfo; -+ tl_submit_infos[idx].pNext = nullptr; -+ vk::SubmitInfo si{ -+ (uint32_t) submission.wait_semaphores.size(), -+ tl_wait_semaphores[idx].data(), -+ stage_flags[idx].data(), -+ 1, -+ &submission.buffer, -+ (uint32_t) submission.signal_semaphores.size(), -+ tl_signal_semaphores[idx].data(), -+ }; -+ si.setPNext(&tl_submit_infos[idx]); -+ submit_infos.push_back(si); -+ } -+ } -+ -+ ctx->q->queue.submit(submit_infos, fence); -+ -+ ctx->seqs.clear(); -+} -+ -+static uint32_t ggml_vk_find_queue_family_index(std::vector& queue_family_props, const vk::QueueFlags& required, const vk::QueueFlags& avoid, int32_t compute_index, uint32_t min_num_queues) { -+ VK_LOG_DEBUG("ggml_vk_find_queue_family_index()"); -+ const uint32_t qfsize = queue_family_props.size(); -+ -+ // Try with avoid preferences first -+ for (uint32_t i = 0; i < qfsize; i++) { -+ if (queue_family_props[i].queueCount >= min_num_queues && (compute_index < 0 || i != (uint32_t) compute_index) && queue_family_props[i].queueFlags & required && !(queue_family_props[i].queueFlags & avoid)) { -+ return i; -+ } -+ } -+ -+ // Fall back to only required -+ for (size_t i = 0; i < qfsize; i++) { -+ if (queue_family_props[i].queueCount >= min_num_queues && (compute_index < 0 || i != (uint32_t) compute_index) && queue_family_props[i].queueFlags & required) { -+ return i; -+ } -+ } -+ -+ // Fall back to reusing compute queue -+ for (size_t i = 0; i < qfsize; i++) { -+ if (queue_family_props[i].queueCount >= min_num_queues && queue_family_props[i].queueFlags & required) { -+ return i; -+ } -+ } -+ -+ // Fall back to ignoring min_num_queries -+ for (size_t i = 0; i < qfsize; i++) { -+ if (queue_family_props[i].queueFlags & required) { -+ return i; -+ } -+ } -+ -+ // All commands that are allowed on a queue that supports transfer operations are also allowed on a queue that supports either graphics or compute operations. -+ // Thus, if the capabilities of a queue family include VK_QUEUE_GRAPHICS_BIT or VK_QUEUE_COMPUTE_BIT, then reporting the VK_QUEUE_TRANSFER_BIT capability separately for that queue family is optional. -+ if (compute_index >= 0) { -+ return compute_index; -+ } -+ -+ std::cerr << "ggml_vulkan: No suitable queue family index found." << std::endl; -+ -+ for(auto &q_family : queue_family_props) { -+ std::cerr << "Queue number: " + std::to_string(q_family.queueCount) << " flags: " + to_string(q_family.queueFlags) << std::endl; -+ } -+ abort(); -+} -+ -+static void ggml_vk_create_queue(vk_device& device, vk_queue& q, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) { -+ VK_LOG_DEBUG("ggml_vk_create_queue()"); -+ std::lock_guard guard(device->mutex); -+ -+ q.queue_family_index = queue_family_index; -+ q.transfer_only = transfer_only; -+ -+ vk::CommandPoolCreateInfo command_pool_create_info_compute(vk::CommandPoolCreateFlags(VK_COMMAND_POOL_CREATE_TRANSIENT_BIT), queue_family_index); -+ q.pool = device->device.createCommandPool(command_pool_create_info_compute); -+ -+ q.cmd_buffer_idx = 0; -+ -+ q.queue = device->device.getQueue(queue_family_index, queue_index); -+ -+ q.stage_flags = stage_flags; -+} -+ -+static vk_context ggml_vk_create_context(ggml_backend_vk_context * ctx, vk_queue& q) { -+ vk_context result = std::make_shared(); -+ VK_LOG_DEBUG("ggml_vk_create_context(" << result << ")"); -+ ctx->gc.contexts.emplace_back(result); -+ result->q = &q; -+ return result; -+} -+ -+static vk_context ggml_vk_create_temporary_context(vk_queue& q) { -+ vk_context result = std::make_shared(); -+ VK_LOG_DEBUG("ggml_vk_create_temporary_context(" << result << ")"); -+ result->q = &q; -+ return result; -+} -+ -+static vk_semaphore * ggml_vk_create_binary_semaphore(ggml_backend_vk_context * ctx) { -+ VK_LOG_DEBUG("ggml_vk_create_timeline_semaphore()"); -+ vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eBinary, 0 }; -+ vk::SemaphoreCreateInfo ci{}; -+ ci.setPNext(&tci); -+ vk::Semaphore semaphore = ctx->device->device.createSemaphore(ci); -+ ctx->gc.semaphores.push_back({ semaphore, 0 }); -+ return &ctx->gc.semaphores[ctx->gc.semaphores.size() - 1]; -+} -+ -+static vk_semaphore * ggml_vk_create_timeline_semaphore(ggml_backend_vk_context * ctx) { -+ VK_LOG_DEBUG("ggml_vk_create_timeline_semaphore()"); -+ if (ctx->semaphore_idx >= ctx->gc.tl_semaphores.size()) { -+ vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eTimeline, 0 }; -+ vk::SemaphoreCreateInfo ci{}; -+ ci.setPNext(&tci); -+ vk::Semaphore semaphore = ctx->device->device.createSemaphore(ci); -+ ctx->gc.tl_semaphores.push_back({ semaphore, 0 }); -+ } -+ return &ctx->gc.tl_semaphores[ctx->semaphore_idx++]; -+} -+ -+static vk::Event ggml_vk_create_event(ggml_backend_vk_context * ctx) { -+ if (ctx->event_idx >= ctx->gc.events.size()) { -+ ctx->gc.events.push_back(ctx->device->device.createEvent({})); -+ } -+ return ctx->gc.events[ctx->event_idx++]; -+} -+ -+static void ggml_vk_queue_cleanup(vk_device& device, vk_queue& q) { -+ VK_LOG_DEBUG("ggml_vk_queue_cleanup()"); -+ std::lock_guard guard(device->mutex); -+ -+ // Requires command buffers to be done -+ device->device.resetCommandPool(q.pool); -+ q.cmd_buffer_idx = 0; -+} -+ -+static uint32_t find_properties(const vk::PhysicalDeviceMemoryProperties* mem_props, vk::MemoryRequirements* mem_req, vk::MemoryPropertyFlags flags) { -+ for (uint32_t i = 0; i < mem_props->memoryTypeCount; ++i) { -+ vk::MemoryType memory_type = mem_props->memoryTypes[i]; -+ if ((mem_req->memoryTypeBits & ((uint64_t)1 << i)) && -+ (flags & memory_type.propertyFlags) == flags && -+ mem_props->memoryHeaps[memory_type.heapIndex].size >= mem_req->size) { -+ return static_cast(i); -+ } -+ } -+ return UINT32_MAX; -+} -+ -+static vk_buffer ggml_vk_create_buffer(vk_device& device, size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags fallback_flags = vk::MemoryPropertyFlags(0)) { -+ VK_LOG_DEBUG("ggml_vk_create_buffer(" << device->name << ", " << size << ", " << to_string(req_flags) << ", " << to_string(fallback_flags) << ")"); -+ if (size > device->max_memory_allocation_size) { -+ throw vk::OutOfDeviceMemoryError("Requested buffer size exceeds device memory allocation limit"); -+ } -+ -+ std::lock_guard guard(device->mutex); -+ -+ vk_buffer buf = std::make_shared(); -+ -+ if (size == 0) { -+ buf->size = 0; -+ return buf; -+ } -+ -+ vk::BufferCreateInfo buffer_create_info{ -+ vk::BufferCreateFlags(), -+ size, -+ vk::BufferUsageFlagBits::eStorageBuffer | vk::BufferUsageFlagBits::eTransferSrc | vk::BufferUsageFlagBits::eTransferDst, -+ vk::SharingMode::eExclusive, -+ 0, -+ nullptr, -+ }; -+ -+ buf->buffer = device->device.createBuffer(buffer_create_info); -+ -+ vk::MemoryRequirements mem_req = device->device.getBufferMemoryRequirements(buf->buffer); -+ -+ vk::PhysicalDeviceMemoryProperties mem_props = device->physical_device.getMemoryProperties(); -+ -+ uint32_t memory_type_index = UINT32_MAX; -+ -+ memory_type_index = find_properties(&mem_props, &mem_req, req_flags); -+ buf->memory_property_flags = req_flags; -+ -+ if (memory_type_index == UINT32_MAX && fallback_flags) { -+ memory_type_index = find_properties(&mem_props, &mem_req, fallback_flags); -+ buf->memory_property_flags = fallback_flags; -+ } -+ -+ if (memory_type_index == UINT32_MAX) { -+ device->device.destroyBuffer(buf->buffer); -+ throw vk::OutOfDeviceMemoryError("No suitable memory type found"); -+ } -+ -+ try { -+ buf->device_memory = device->device.allocateMemory({ mem_req.size, memory_type_index }); -+ } catch (const vk::SystemError& e) { -+ if (buf->memory_property_flags != fallback_flags) { -+ // Try again with fallback flags -+ memory_type_index = find_properties(&mem_props, &mem_req, fallback_flags); -+ buf->memory_property_flags = fallback_flags; -+ -+ try { -+ buf->device_memory = device->device.allocateMemory({ mem_req.size, memory_type_index }); -+ } -+ catch (const vk::SystemError& e) { -+ device->device.destroyBuffer(buf->buffer); -+ throw e; -+ } -+ } else { -+ // Out of Host/Device memory, clean up buffer -+ device->device.destroyBuffer(buf->buffer); -+ throw e; -+ } -+ } -+ buf->ptr = nullptr; -+ -+ if (buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { -+ buf->ptr = device->device.mapMemory(buf->device_memory, 0, VK_WHOLE_SIZE); -+ } -+ -+ device->device.bindBufferMemory(buf->buffer, buf->device_memory, 0); -+ -+ buf->device = device; -+ buf->size = size; -+ -+#ifdef GGML_VULKAN_MEMORY_DEBUG -+ device->memory_logger->log_allocation(buf, size); -+#endif -+ -+ return buf; -+} -+ -+static vk_buffer ggml_vk_create_buffer_check(vk_device& device, size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags fallback_flags = vk::MemoryPropertyFlags(0)) { -+ try { -+ return ggml_vk_create_buffer(device, size, req_flags, fallback_flags); -+ } catch (const vk::SystemError& e) { -+ std::cerr << "ggml_vulkan: Memory allocation of size " << size << " failed." << std::endl; -+ std::cerr << "ggml_vulkan: " << e.what() << std::endl; -+ throw e; -+ } -+} -+ -+static vk_buffer ggml_vk_create_buffer_device(vk_device& device, size_t size) { -+ vk_buffer buf; -+ try { -+ if (device->uma) { -+ // Fall back to host memory type -+ buf = ggml_vk_create_buffer(device, size, vk::MemoryPropertyFlagBits::eDeviceLocal, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); -+ } else { -+ // use rebar if available, otherwise fallback to device only visible memory -+ buf = ggml_vk_create_buffer(device, size, vk::MemoryPropertyFlagBits::eDeviceLocal | vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent, vk::MemoryPropertyFlagBits::eDeviceLocal); -+ } -+ } catch (const vk::SystemError& e) { -+ std::cerr << "ggml_vulkan: Device memory allocation of size " << size << " failed." << std::endl; -+ std::cerr << "ggml_vulkan: " << e.what() << std::endl; -+ throw e; -+ } -+ -+ return buf; -+} -+ -+static void ggml_vk_destroy_buffer(vk_buffer& buf) { -+ if (buf == nullptr) { -+ return; -+ } -+ -+#ifdef GGML_VULKAN_MEMORY_DEBUG -+ if (buf->device != nullptr) { -+ buf->device->memory_logger->log_deallocation(buf); -+ } -+#endif -+ -+ buf.reset(); -+} -+ -+static vk_subbuffer ggml_vk_subbuffer(vk_buffer& buf) { -+ return { buf, 0, VK_WHOLE_SIZE }; -+} -+ -+static void ggml_vk_sync_buffers(vk_context& ctx) { -+ VK_LOG_DEBUG("ggml_vk_sync_buffers()"); -+ -+ const bool transfer_queue = ctx->q->transfer_only; -+ -+ ctx->s->buffer.pipelineBarrier( -+ ctx->q->stage_flags, -+ ctx->q->stage_flags, -+ {}, -+ { { -+ { !transfer_queue ? (vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) : (vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) }, -+ { !transfer_queue ? (vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) : (vk::AccessFlagBits::eTransferRead | vk::AccessFlagBits::eTransferWrite) } -+ } }, -+ {}, -+ {} -+ ); -+} -+ -+static void ggml_vk_wait_events(vk_context& ctx, std::vector&& events) { -+ VK_LOG_DEBUG("ggml_vk_wait_events()"); -+ if (events.empty()) { -+ return; -+ } -+ -+ ctx->s->buffer.waitEvents( -+ events, -+ ctx->q->stage_flags, -+ ctx->q->stage_flags, -+ {}, -+ {}, -+ {} -+ ); -+} -+ -+// number of rows/cols for flash attention shader -+static constexpr uint32_t flash_attention_num_small_rows = 32; -+static std::array fa_rows_cols(uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) { -+ GGML_UNUSED(clamp); -+ -+ // small rows, large cols -+ if (small_rows) { -+ return {flash_attention_num_small_rows, 128}; -+ } -+ // small cols to reduce register count -+ if (ggml_is_quantized(type) || D == 256) { -+ return {64, 32}; -+ } -+ return {64, 64}; -+}; -+ -+static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vector& warptile, bool mul_mat_id) { -+ // Needs to be kept up to date on shader changes -+ const uint32_t bank_conflict_offset = device->coopmat_support ? 8 : 1; -+ const uint32_t type_size = device->fp16 ? sizeof(ggml_fp16_t) : sizeof(float); -+ const uint32_t warps = warptile[0] / warptile[10]; -+ -+ const uint32_t load_bufs = (warptile[1] + warptile[2]) * (warptile[3] + bank_conflict_offset) * type_size; -+ const uint32_t mmid_row_ids = mul_mat_id ? 3072 * sizeof(uint32_t) : 0; -+ const uint32_t coopmat_stage = device->coopmat_support ? warptile[7] * warptile[8] / warps * sizeof(float) : 0; -+ -+ return (load_bufs + mmid_row_ids + coopmat_stage) <= device->properties.limits.maxComputeSharedMemorySize; -+} -+ -+static void ggml_vk_load_shaders(vk_device& device) { -+ VK_LOG_DEBUG("ggml_vk_load_shaders(" << device->name << ")"); -+ -+ std::cerr << "ggml_vulkan: Compiling shaders"; -+ -+ // some shaders have a minimum subgroup size -+ const uint32_t subgroup_size_16 = std::max(device->subgroup_size, 16u); -+ const uint32_t subgroup_size_32 = std::max(device->subgroup_size, 32u); -+ -+ // mulmat -+ std::vector l_warptile, m_warptile, s_warptile, -+ l_warptile_mmq, m_warptile_mmq, s_warptile_mmq, -+ l_warptile_mmq_k, m_warptile_mmq_k, s_warptile_mmq_k, -+ l_warptile_mmqid, m_warptile_mmqid, s_warptile_mmqid; -+ std::array l_wg_denoms, m_wg_denoms, s_wg_denoms, -+ l_mmq_wg_denoms, m_mmq_wg_denoms, s_mmq_wg_denoms, -+ l_mmq_wg_denoms_k, m_mmq_wg_denoms_k, s_mmq_wg_denoms_k, -+ l_mmqid_wg_denoms, m_mmqid_wg_denoms, s_mmqid_wg_denoms; -+ -+ uint32_t l_align, m_align, s_align; -+ if (device->coopmat2) { -+ // spec constants and tile sizes for non-quant matmul/matmul_id -+ l_warptile = { 256, 128, 256, 64 }; -+ m_warptile = { 256, 128, 128, 64 }; -+ s_warptile = { 128, 64, 64, 64 }; -+ l_wg_denoms = {128, 256, 1 }; -+ m_wg_denoms = {128, 128, 1 }; -+ s_wg_denoms = { 64, 64, 1 }; -+ -+ // spec constants and tile sizes for quant matmul (non-Qi_K) -+ l_warptile_mmq = { 256, 128, 256, 64 }; -+ m_warptile_mmq = { 256, 128, 128, 64 }; -+ s_warptile_mmq = { 256, 128, 128, 64 }; -+ l_mmq_wg_denoms = { 128, 256, 1 }; -+ m_mmq_wg_denoms = { 128, 128, 1 }; -+ s_mmq_wg_denoms = { 128, 128, 1 }; -+ -+ // spec constants and tile sizes for quant matmul (Qi_K) -+ l_warptile_mmq_k = { 256, 128, 512, 16 }; -+ m_warptile_mmq_k = { 256, 128, 256, 16 }; -+ s_warptile_mmq_k = { 256, 32, 128, 64 }; -+ l_mmq_wg_denoms_k = { 128, 512, 1 }; -+ m_mmq_wg_denoms_k = { 128, 256, 1 }; -+ s_mmq_wg_denoms_k = { 32, 128, 1 }; -+ -+ // spec constants and tile sizes for quant matmul_id -+ l_warptile_mmqid = { 256, 128, 128, 16 }; -+ m_warptile_mmqid = { 256, 128, 64, 16 }; -+ s_warptile_mmqid = { 256, 64, 64, 16 }; -+ l_mmqid_wg_denoms = { 128, 128, 1 }; -+ m_mmqid_wg_denoms = { 128, 64, 1 }; -+ s_mmqid_wg_denoms = { 64, 64, 1 }; -+ -+ l_align = 128; -+ m_align = 64; -+ s_align = 32; -+ } else { -+ // Matrix cores require different warp group sizes -+ const uint32_t tm_l = device->coopmat_support ? device->coopmat_m : 4; -+ const uint32_t tm_m = device->coopmat_support ? device->coopmat_m : 4; -+ const uint32_t tm_s = device->coopmat_support ? device->coopmat_m : 2; -+ const uint32_t tn_l = device->coopmat_support ? device->coopmat_n : 4; -+ const uint32_t tn_m = device->coopmat_support ? device->coopmat_n : 2; -+ const uint32_t tn_s = device->coopmat_support ? device->coopmat_n : 2; -+ const uint32_t tk_l = device->coopmat_support ? device->coopmat_k : 1; -+ const uint32_t tk_m = device->coopmat_support ? device->coopmat_k : 1; -+ const uint32_t tk_s = device->coopmat_support ? device->coopmat_k : 1; -+ -+ l_warptile = { 128, 128, 128, 16, device->subgroup_size * 2, 64, 2, tm_l, tn_l, tk_l, device->subgroup_size }; -+ m_warptile = { 128, 64, 64, 16, device->subgroup_size, 32, 2, tm_m, tn_m, tk_m, device->subgroup_size }; -+ s_warptile = { subgroup_size_16, 32, 32, 16, 32, 32, 2, tm_s, tn_s, tk_s, device->subgroup_size }; -+ -+ l_warptile_mmq = { 128, 128, 128, 32, device->subgroup_size * 2, 64, 2, tm_l, tn_l, tk_l, device->subgroup_size }; -+ m_warptile_mmq = { 128, 64, 64, 32, device->subgroup_size, 32, 2, tm_m, tn_m, tk_m, device->subgroup_size }; -+ s_warptile_mmq = { subgroup_size_32, 32, 32, 32, 32, 32, 2, tm_s, tn_s, tk_s, device->subgroup_size }; -+ -+ l_mmq_wg_denoms = l_wg_denoms = {128, 128, 1 }; -+ m_mmq_wg_denoms = m_wg_denoms = { 64, 64, 1 }; -+ s_mmq_wg_denoms = s_wg_denoms = { 32, 32, 1 }; -+ l_align = 128; -+ m_align = 64; -+ s_align = 32; -+ -+ // Fallback to smaller sizes if there's not enough shared memory. Given the current shaders -+ // and tile sizes, this should handle 16KB, 32KB, and 48KB+. -+ // This logic doesn't explicitly account for the 12KB row_ids in the mul_mat_mat_id shaders. -+ // But the numbers happen to work out for 32KB shared memory size that when using the medium -+ // size there's enough room for everything, and we assert for this. -+ uint32_t shmem_needed = (l_warptile[1] + l_warptile[2]) * (l_warptile[3] + 1) * sizeof(float); -+ if (shmem_needed > device->properties.limits.maxComputeSharedMemorySize) { -+ l_warptile = m_warptile; -+ l_wg_denoms = m_wg_denoms; -+ shmem_needed = (l_warptile[1] + l_warptile[2]) * (l_warptile[3] + 1) * sizeof(float); -+ GGML_ASSERT(shmem_needed <= device->properties.limits.maxComputeSharedMemorySize); -+ } -+ if (device->properties.limits.maxComputeSharedMemorySize >= 32768) { -+ // assert mul_mat_mat_id shaders will fit. -+ GGML_ASSERT(shmem_needed + 3072*4 <= device->properties.limits.maxComputeSharedMemorySize); -+ } -+ -+ shmem_needed = (l_warptile_mmq[1] + l_warptile_mmq[2]) * (l_warptile_mmq[3] + 1) * sizeof(float); -+ if (shmem_needed > device->properties.limits.maxComputeSharedMemorySize) { -+ if (device->properties.limits.maxComputeSharedMemorySize == 32768) { -+ l_warptile_mmq = m_warptile_mmq; -+ l_mmq_wg_denoms = m_mmq_wg_denoms; -+ } else { -+ l_warptile_mmq = s_warptile_mmq; -+ l_mmq_wg_denoms = s_mmq_wg_denoms; -+ } -+ shmem_needed = (l_warptile_mmq[1] + l_warptile_mmq[2]) * (l_warptile_mmq[3] + 1) * sizeof(float); -+ GGML_ASSERT(shmem_needed <= device->properties.limits.maxComputeSharedMemorySize); -+ } -+ if (device->properties.limits.maxComputeSharedMemorySize >= 32768) { -+ // assert mul_mat_mat_id shaders will fit. -+ GGML_ASSERT(shmem_needed + 3072*4 <= device->properties.limits.maxComputeSharedMemorySize); -+ } -+ // Disable medium and large matrix multiplication if not enough shared memory is available -+ // Check mmq warptiles as the largest configuration -+ // Throw an error if not enough for any matrix multiplication is available -+ if (!ggml_vk_matmul_shmem_support(device, s_warptile_mmq, false)) { -+ std::cerr << "ggml_vulkan: Error: Shared memory size too small for matrix multiplication." << std::endl; -+ throw std::runtime_error("Shared memory size too small for matrix multiplication."); -+ } else if (!ggml_vk_matmul_shmem_support(device, m_warptile_mmq, false)) { -+ device->mul_mat_m = false; -+ device->mul_mat_l = false; -+ } else if (!ggml_vk_matmul_shmem_support(device, l_warptile_mmq, false)) { -+ device->mul_mat_l = false; -+ } -+ -+ // Disable mul_mat_id if not enough shared memory is available -+ if (!ggml_vk_matmul_shmem_support(device, s_warptile_mmq, true)) { -+ device->mul_mat_id_s = false; -+ device->mul_mat_id_m = false; -+ device->mul_mat_id_l = false; -+ } else if (!ggml_vk_matmul_shmem_support(device, m_warptile_mmq, true)) { -+ device->mul_mat_id_m = false; -+ device->mul_mat_id_l = false; -+ } else if (!ggml_vk_matmul_shmem_support(device, l_warptile_mmq, true)) { -+ device->mul_mat_id_l = false; -+ } -+ } -+ -+ device->pipeline_matmul_f32 = std::make_shared(); -+ device->pipeline_matmul_f32_f16 = std::make_shared(); -+ -+ device->pipeline_matmul_id_f32 = std::make_shared(); -+ -+ std::vector> compiles; -+ auto const &ggml_vk_create_pipeline = [&](vk_device& device, vk_pipeline& pipeline, const std::string &name, size_t spv_size, const void* spv_data, const std::string &entrypoint, -+ uint32_t parameter_count, uint32_t push_constant_size, std::array wg_denoms, const std::vector& specialization_constants, -+ uint32_t align, bool disable_robustness = false, bool require_full_subgroups = false, uint32_t required_subgroup_size = 0) { -+ { -+ // wait until fewer than N compiles are in progress -+ uint32_t N = std::max(1u, std::thread::hardware_concurrency()); -+ std::unique_lock guard(compile_count_mutex); -+ while (compile_count >= N) { -+ compile_count_cond.wait(guard); -+ } -+ compile_count++; -+ } -+ compiles.push_back(std::async(ggml_vk_create_pipeline_func, std::ref(device), std::ref(pipeline), name, spv_size, spv_data, entrypoint, -+ parameter_count, push_constant_size, wg_denoms, specialization_constants, align, disable_robustness, require_full_subgroups, required_subgroup_size)); -+ }; -+ -+#if defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) -+ if (device->coopmat2) { -+ -+ auto const &fa_wg_denoms = [&](uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) -> std::array { -+ return {fa_rows_cols(D, clamp, type, small_rows)[0], 1, 1}; -+ }; -+ -+ auto const &fa_spec_constants = [&](uint32_t D, uint32_t clamp, ggml_type type, bool small_rows) -> std::vector { -+ // For large number of rows, 128 invocations seems to work best. -+ // For small number of rows (e.g. N==1), 256 works better. But matrix granularity for 256 is 32, so we -+ // can't use 256 for D==80. -+ uint32_t wg_size = (small_rows && (D % 32) == 0) ? 256 : 128; -+ auto rows_cols = fa_rows_cols(D, clamp, type, small_rows); -+ return {wg_size, rows_cols[0], rows_cols[1], (D), clamp}; -+ }; -+ -+#define CREATE_FA2(TYPE, NAMELC, D) \ -+ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][0][0], "flash_attn_f32_f16_D" #D "_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,false), fa_spec_constants(D,1,TYPE,false), 1); \ -+ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][0][1], "flash_attn_f32_f16_D" #D "_aligned_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,false), fa_spec_constants(D,0,TYPE,false), fa_rows_cols(D,0,TYPE,false)[1]); \ -+ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][0][0], "flash_attn_f32_f16_D" #D "_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,false), fa_spec_constants(D,1,TYPE,false), 1); \ -+ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][0][1], "flash_attn_f32_f16_D" #D "_aligned_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,false), fa_spec_constants(D,0,TYPE,false), fa_rows_cols(D,0,TYPE,false)[1]); \ -+ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][1][0], "flash_attn_f32_f16_D" #D "_f16acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,true), fa_spec_constants(D,1,TYPE,true), 1); \ -+ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][0][1][1], "flash_attn_f32_f16_D" #D "_aligned_f16acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_len, flash_attn_f32_f16_ ## NAMELC ## _f16acc_cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,true), fa_spec_constants(D,0,TYPE,true), fa_rows_cols(D,0,TYPE,true)[1]); \ -+ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][1][0], "flash_attn_f32_f16_D" #D "_f32acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,1,TYPE,true), fa_spec_constants(D,1,TYPE,true), 1); \ -+ ggml_vk_create_pipeline(device, device->pipeline_flash_attn_f32_f16_D ## D[TYPE][1][1][1], "flash_attn_f32_f16_D" #D "_aligned_f32acc_smallrows" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _cm2_len, flash_attn_f32_f16_ ## NAMELC ## _cm2_data, "main", 5, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(D,0,TYPE,true), fa_spec_constants(D,0,TYPE,true), fa_rows_cols(D,0,TYPE,true)[1]); \ -+ -+#define CREATE_FA(TYPE, NAMELC) \ -+ CREATE_FA2(TYPE, NAMELC, 64) \ -+ CREATE_FA2(TYPE, NAMELC, 80) \ -+ CREATE_FA2(TYPE, NAMELC, 96) \ -+ CREATE_FA2(TYPE, NAMELC, 112) \ -+ CREATE_FA2(TYPE, NAMELC, 128) \ -+ CREATE_FA2(TYPE, NAMELC, 256) -+ -+ CREATE_FA(GGML_TYPE_F16, f16) -+ CREATE_FA(GGML_TYPE_Q4_0, q4_0) -+ CREATE_FA(GGML_TYPE_Q4_1, q4_1) -+ CREATE_FA(GGML_TYPE_Q5_0, q5_0) -+ CREATE_FA(GGML_TYPE_Q5_1, q5_1) -+ CREATE_FA(GGML_TYPE_Q8_0, q8_0) -+ // K dequants currently disabled because D dimension is rounded up to 256 and runs inefficiently -+ //CREATE_FA(GGML_TYPE_Q2_K, q2_k) -+ //CREATE_FA(GGML_TYPE_Q3_K, q3_k) -+ //CREATE_FA(GGML_TYPE_Q4_K, q4_k) -+ //CREATE_FA(GGML_TYPE_Q5_K, q5_k) -+ //CREATE_FA(GGML_TYPE_Q6_K, q6_k) -+ CREATE_FA(GGML_TYPE_IQ4_NL, iq4_nl) -+#undef CREATE_FA -+ -+ // Create 6 variants, {s,m,l}x{unaligned,aligned} -+#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _cm2_len, NAMELC ## _aligned ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ -+ -+ // Create 2 variants, {f16,f32} accumulator -+#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ -+ CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ -+ CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ -+ -+ CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3) -+ -+ CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3) -+ CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f16, _f16acc, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f16, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) -+ -+ CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_id_push_constants, 4) -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, , mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 4) -+#undef CREATE_MM -+#undef CREATE_MM2 -+ } else -+#endif // defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) -+ if (device->coopmat_support) { -+ // Create 6 variants, {s,m,l}x{unaligned,aligned} -+#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ -+ if (device->mul_mat ## ID ## _l) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1, false, true); \ -+ if (device->mul_mat ## ID ## _m) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1, false, true); \ -+ if (device->mul_mat ## ID ## _s) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _coopmat_len, NAMELC ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1, false, true); \ -+ if (device->mul_mat ## ID ## _l) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align, false, true); \ -+ if (device->mul_mat ## ID ## _m) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align, false, true); \ -+ if (device->mul_mat ## ID ## _s) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _coopmat_len, NAMELC ## _aligned ## F16ACC ## _coopmat_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align, false, true); \ -+ -+ // Create 2 variants, {f16,f32} accumulator -+#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ -+ if (device->coopmat_acc_f16_support) { \ -+ CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ -+ } \ -+ if (device->coopmat_acc_f32_support) { \ -+ CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ -+ } \ -+ -+ CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ -+ if (device->coopmat_acc_f16_support) { -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ } else { -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, , wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ } -+ -+ // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. -+ if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { -+ CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); -+ CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); -+ CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); -+ -+ if (device->coopmat_acc_f16_support) { -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, _f16acc, wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ } else { -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, , wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ } -+ } -+#undef CREATE_MM2 -+#undef CREATE_MM -+ } else if (device->fp16) { -+ // Create 6 variants, {s,m,l}x{unaligned,aligned} -+#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ -+ if (device->mul_mat ## ID ## _l) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ -+ if (device->mul_mat ## ID ## _m) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ -+ if (device->mul_mat ## ID ## _s) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ -+ if (device->mul_mat ## ID ## _l) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ -+ if (device->mul_mat ## ID ## _m) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ -+ if (device->mul_mat ## ID ## _s) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _len, NAMELC ## _aligned ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ -+ -+ // Create 2 variants, {f16,f32} accumulator -+#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ -+ CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ -+ CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ -+ -+ CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM2(pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f16acc, matmul_q4_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f16acc, matmul_q5_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f16acc, matmul_q5_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f16acc, matmul_q2_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f16acc, matmul_q3_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f16acc, matmul_q4_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f16acc, matmul_q5_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f16acc, matmul_q6_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc, matmul_iq4_nl_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ -+ // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. -+ if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { -+ CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); -+ CREATE_MM2(pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); -+ CREATE_MM2(pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f16acc, matmul_id_q4_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f16acc, matmul_id_q4_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f16acc, matmul_id_q5_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f16acc, matmul_id_q5_1_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f16acc, matmul_id_q8_0_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f16acc, matmul_id_q2_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f16acc, matmul_id_q3_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f16acc, matmul_id_q4_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f16acc, matmul_id_q5_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f16acc, matmul_id_q6_k_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f16acc, matmul_id_iq4_nl_f32, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ } -+#undef CREATE_MM2 -+#undef CREATE_MM -+ } else { -+ // Create 6 variants, {s,m,l}x{unaligned,aligned} -+#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ -+ if (device->mul_mat ## ID ## _l) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ -+ if (device->mul_mat ## ID ## _m) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ -+ if (device->mul_mat ## ID ## _s) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ -+ if (device->mul_mat ## ID ## _l) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align); \ -+ if (device->mul_mat ## ID ## _m) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align); \ -+ if (device->mul_mat ## ID ## _s) \ -+ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## _aligned ## F16ACC ## _fp32_len, NAMELC ## _aligned ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align); \ -+ -+ CREATE_MM(pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_matmul_f16.f32acc, matmul_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_matmul_f16_f32.f32acc, matmul_f16_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f32acc, matmul_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f32acc, matmul_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f32acc, matmul_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f32acc, matmul_q5_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f32acc, matmul_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f32acc, matmul_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f32acc, matmul_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f32acc, matmul_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f32acc, matmul_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f32acc, matmul_q6_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ CREATE_MM(pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f32acc, matmul_iq4_nl_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); -+ -+ // If there's not enough shared memory for row_ids and the result tile, don't create these pipelines. -+ if (device->mul_mat_id_s || device->mul_mat_id_m || device->mul_mat_id_l) { -+ CREATE_MM(pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); -+ CREATE_MM(pipeline_matmul_id_f16.f32acc, matmul_id_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); -+ CREATE_MM(pipeline_matmul_id_f16_f32.f32acc, matmul_id_f16_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f32acc, matmul_id_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f32acc, matmul_id_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f32acc, matmul_id_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f32acc, matmul_id_q5_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f32acc, matmul_id_q6_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ CREATE_MM(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f32acc, matmul_id_iq4_nl_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, 4, _id); -+ } -+#undef CREATE_MM -+ } -+ -+ // mul mat vec -+ -+ // the number of rows computed per shader depends on GPU model and quant -+ uint32_t rm_stdq = 1; -+ uint32_t rm_kq = 2; -+ if (device->vendor_id == VK_VENDOR_ID_AMD) { -+ if (device->subgroup_min_size == 64 && device->subgroup_max_size == 64) { // GCN -+ rm_stdq = 2; -+ rm_kq = 4; -+ } -+ } else if (device->vendor_id == VK_VENDOR_ID_INTEL) -+ rm_stdq = 2; -+ -+ for (uint32_t i = 0; i < mul_mat_vec_max_cols; ++i) { -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_F32 ][i], "mul_mat_vec_f32_f32_f32_"+std::to_string(i+1), mul_mat_vec_f32_f32_f32_len, mul_mat_vec_f32_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f32_f32_"+std::to_string(i+1), mul_mat_vec_f16_f32_f32_len, mul_mat_vec_f16_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_0_f32_f32_len, mul_mat_vec_q4_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_1_f32_f32_len, mul_mat_vec_q4_1_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_0_f32_f32_len, mul_mat_vec_q5_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_1_f32_f32_len, mul_mat_vec_q5_1_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f32_f32_"+std::to_string(i+1), mul_mat_vec_q8_0_f32_f32_len, mul_mat_vec_q8_0_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q2_k_f32_f32_len, mul_mat_vec_q2_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q3_k_f32_f32_len, mul_mat_vec_q3_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q4_k_f32_f32_len, mul_mat_vec_q4_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q5_k_f32_f32_len, mul_mat_vec_q5_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_f32_f32_"+std::to_string(i+1), mul_mat_vec_q6_k_f32_f32_len, mul_mat_vec_q6_k_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f32_f32_"+std::to_string(i+1), mul_mat_vec_iq4_nl_f32_f32_len, mul_mat_vec_iq4_nl_f32_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq, i+1}, 1, true); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_F32 ][i], "mul_mat_vec_f32_f16_f32_"+std::to_string(i+1), mul_mat_vec_f32_f16_f32_len, mul_mat_vec_f32_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f16_f32_"+std::to_string(i+1), mul_mat_vec_f16_f16_f32_len, mul_mat_vec_f16_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {device->subgroup_size, 2, i+1}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_0_f16_f32_len, mul_mat_vec_q4_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_1_f16_f32_len, mul_mat_vec_q4_1_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_0_f16_f32_len, mul_mat_vec_q5_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_1_f16_f32_len, mul_mat_vec_q5_1_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f16_f32_"+std::to_string(i+1), mul_mat_vec_q8_0_f16_f32_len, mul_mat_vec_q8_0_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q2_k_f16_f32_len, mul_mat_vec_q2_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q3_k_f16_f32_len, mul_mat_vec_q3_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q4_k_f16_f32_len, mul_mat_vec_q4_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q5_k_f16_f32_len, mul_mat_vec_q5_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_f16_f32_"+std::to_string(i+1), mul_mat_vec_q6_k_f16_f32_len, mul_mat_vec_q6_k_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq, i+1}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f16_f32_"+std::to_string(i+1), mul_mat_vec_iq4_nl_f16_f32_len, mul_mat_vec_iq4_nl_f16_f32_data, "main", 3, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq, i+1}, 1, true); -+ } -+ -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_F32 ], "mul_mat_vec_id_f32_f32", mul_mat_vec_id_f32_f32_len, mul_mat_vec_id_f32_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {device->subgroup_size, 2}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_F16 ], "mul_mat_vec_id_f16_f32", mul_mat_vec_id_f16_f32_len, mul_mat_vec_id_f16_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {device->subgroup_size, 2}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_0], "mul_mat_vec_id_q4_0_f32", mul_mat_vec_id_q4_0_f32_len, mul_mat_vec_id_q4_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_1], "mul_mat_vec_id_q4_1_f32", mul_mat_vec_id_q4_1_f32_len, mul_mat_vec_id_q4_1_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_0], "mul_mat_vec_id_q5_0_f32", mul_mat_vec_id_q5_0_f32_len, mul_mat_vec_id_q5_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_1], "mul_mat_vec_id_q5_1_f32", mul_mat_vec_id_q5_1_f32_len, mul_mat_vec_id_q5_1_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {device->subgroup_size, 2*rm_stdq}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q8_0], "mul_mat_vec_id_q8_0_f32", mul_mat_vec_id_q8_0_f32_len, mul_mat_vec_id_q8_0_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq, 1, 1}, {device->subgroup_size, 1*rm_stdq}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q2_K], "mul_mat_vec_id_q2_k_f32", mul_mat_vec_id_q2_k_f32_len, mul_mat_vec_id_q2_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q3_K], "mul_mat_vec_id_q3_k_f32", mul_mat_vec_id_q3_k_f32_len, mul_mat_vec_id_q3_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q4_K], "mul_mat_vec_id_q4_k_f32", mul_mat_vec_id_q4_k_f32_len, mul_mat_vec_id_q4_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q5_K], "mul_mat_vec_id_q5_k_f32", mul_mat_vec_id_q5_k_f32_len, mul_mat_vec_id_q5_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_Q6_K], "mul_mat_vec_id_q6_k_f32", mul_mat_vec_id_q6_k_f32_len, mul_mat_vec_id_q6_k_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {subgroup_size_16, rm_kq}, 1, true); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[GGML_TYPE_IQ4_NL], "mul_mat_vec_id_iq4_nl_f32", mul_mat_vec_id_iq4_nl_f32_len, mul_mat_vec_id_iq4_nl_f32_data, "main", 4, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {subgroup_size_16, 2*rm_stdq}, 1, true); -+ -+ // dequant shaders -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_F32 ], "f32_to_f16", dequant_f32_len, dequant_f32_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_0], "dequant_q4_0", dequant_q4_0_len, dequant_q4_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_1], "dequant_q4_1", dequant_q4_1_len, dequant_q4_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_1], "dequant_q5_1", dequant_q5_1_len, dequant_q5_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_K], "dequant_q4_k", dequant_q4_k_len, dequant_q4_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_K], "dequant_q5_k", dequant_q5_k_len, dequant_q5_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q6_K], "dequant_q6_k", dequant_q6_k_len, dequant_q6_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ4_NL], "dequant_iq4_nl", dequant_iq4_nl_len, dequant_iq4_nl_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); -+ -+ // get_rows -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_F32 ], "get_rows_f32", get_rows_f32_len, get_rows_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_F16 ], "get_rows_f16", get_rows_f16_len, get_rows_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_0], "get_rows_q4_0", get_rows_q4_0_len, get_rows_q4_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_1], "get_rows_q4_1", get_rows_q4_1_len, get_rows_q4_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_0], "get_rows_q5_0", get_rows_q5_0_len, get_rows_q5_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_1], "get_rows_q5_1", get_rows_q5_1_len, get_rows_q5_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q8_0], "get_rows_q8_0", get_rows_q8_0_len, get_rows_q8_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_IQ4_NL], "get_rows_iq4_nl", get_rows_iq4_nl_len, get_rows_iq4_nl_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_F32 ], "get_rows_f32_f32", get_rows_f32_f32_len, get_rows_f32_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_F16 ], "get_rows_f16_f32", get_rows_f16_f32_len, get_rows_f16_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_0], "get_rows_q4_0_f32", get_rows_q4_0_f32_len, get_rows_q4_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_1], "get_rows_q4_1_f32", get_rows_q4_1_f32_len, get_rows_q4_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_0], "get_rows_q5_0_f32", get_rows_q5_0_f32_len, get_rows_q5_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_1], "get_rows_q5_1_f32", get_rows_q5_1_f32_len, get_rows_q5_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q8_0], "get_rows_q8_0_f32", get_rows_q8_0_f32_len, get_rows_q8_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_IQ4_NL], "get_rows_iq4_nl_f32", get_rows_iq4_nl_f32_len, get_rows_iq4_nl_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_matmul_split_k_reduce, "split_k_reduce", split_k_reduce_len, split_k_reduce_data, "main", 2, 2 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_mul_mat_vec_p021_f16_f32, "mul_mat_vec_p021_f16_f32", mul_mat_vec_p021_f16_f32_len, mul_mat_vec_p021_f16_f32_data, "main", 3, 6 * sizeof(uint32_t), {1, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_mul_mat_vec_nc_f16_f32, "mul_mat_vec_nc_f16_f32", mul_mat_vec_nc_f16_f32_len, mul_mat_vec_nc_f16_f32_data, "main", 3, 7 * sizeof(uint32_t), {1, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_norm_f32, "norm_f32", norm_f32_len, norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_group_norm_f32, "group_norm_f32", group_norm_f32_len, group_norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_rms_norm_f32, "rms_norm_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f32, "cpy_f32_f32", cpy_f32_f32_len, cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f16, "cpy_f32_f16", cpy_f32_f16_len, cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_cpy_f16_f16, "cpy_f16_f16", cpy_f16_f16_len, cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f32, "contig_cpy_f32_f32", contig_cpy_f32_f32_len, contig_cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f16, "contig_cpy_f32_f16", contig_cpy_f32_f16_len, contig_cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f16_f16, "contig_cpy_f16_f16", contig_cpy_f16_f16_len, contig_cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_add_f32, "add_f32", add_f32_len, add_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_add_f32_norepeat, "add_f32_norepeat", add_f32_len, add_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_add_f16_f32_f16, "add_f16_f32_f16", add_f16_f32_f16_len, add_f16_f32_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_add_f16_f32_f16_norepeat, "add_f16_f32_f16_norepeat", add_f16_f32_f16_len, add_f16_f32_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_acc_f32, "acc_f32", acc_f32_len, acc_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_mul_f32, "mul_f32", mul_f32_len, mul_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_mul_f32_norepeat, "mul_f32_norepeat", mul_f32_len, mul_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_div_f32, "div_f32", div_f32_len, div_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {0}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_div_f32_norepeat, "div_f32_norepeat", div_f32_len, div_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_concat_f32, "concat_f32", concat_f32_len, concat_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_concat_f16, "concat_f16", concat_f16_len, concat_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_concat_i32, "concat_i32", concat_i32_len, concat_i32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_upscale_f32, "upscale_f32", upscale_f32_len, upscale_f32_data, "main", 2, sizeof(vk_op_upscale_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_scale_f32, "scale_f32", scale_f32_len, scale_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_sqr_f32, "sqr_f32", sqr_f32_len, sqr_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_sin_f32, "sin_f32", sin_f32_len, sin_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_cos_f32, "cos_f32", cos_f32_len, cos_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_clamp_f32, "clamp_f32", clamp_f32_len, clamp_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_pad_f32, "pad_f32", pad_f32_len, pad_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_repeat_f32, "repeat_f32", repeat_f32_len, repeat_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_gelu_f32, "gelu_f32", gelu_f32_len, gelu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_gelu_quick_f32, "gelu_quick_f32", gelu_quick_f32_len, gelu_quick_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_silu_f32, "silu_f32", silu_f32_len, silu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_relu_f32, "relu_f32", relu_f32_len, relu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_leaky_relu_f32, "leaky_relu_f32", leaky_relu_f32_len, leaky_relu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_tanh_f32, "tanh_f32", tanh_f32_len, tanh_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_diag_mask_inf_f32, "diag_mask_inf_f32", diag_mask_inf_f32_len, diag_mask_inf_f32_data, "main", 2, sizeof(vk_op_diag_mask_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32, "soft_max_f32", soft_max_f32_len, soft_max_f32_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_wg512, "soft_max_f32_wg512", soft_max_f32_len, soft_max_f32_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 512 }, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_f16, "soft_max_f32_f16", soft_max_f32_f16_len, soft_max_f32_f16_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_f16_wg512, "soft_max_f32_f16_wg512", soft_max_f32_f16_len, soft_max_f32_f16_data, "main", 3, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 512 }, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f32, "rope_norm_f32", rope_norm_f32_len, rope_norm_f32_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f32, "rope_neox_f32", rope_neox_f32_len, rope_neox_f32_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); -+ -+ if (device->float_controls_rte_fp16) { -+ ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f16, "rope_norm_f16", rope_norm_f16_rte_len, rope_norm_f16_rte_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f16, "rope_neox_f16", rope_neox_f16_rte_len, rope_neox_f16_rte_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); -+ } else { -+ ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f16, "rope_norm_f16", rope_norm_f16_len, rope_norm_f16_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); -+ ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f16, "rope_neox_f16", rope_neox_f16_len, rope_neox_f16_data, "main", 4, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1); -+ } -+ -+ ggml_vk_create_pipeline(device, device->pipeline_argsort_f32, "argsort_f32", argsort_f32_len, argsort_f32_data, "main", 2, sizeof(vk_op_argsort_push_constants), {1024, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_im2col_f32, "im2col_f32", im2col_f32_len, im2col_f32_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); -+ if (device->float_controls_rte_fp16) { -+ ggml_vk_create_pipeline(device, device->pipeline_im2col_f32_f16, "im2col_f32_f16", im2col_f32_f16_rte_len, im2col_f32_f16_rte_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); -+ } else { -+ ggml_vk_create_pipeline(device, device->pipeline_im2col_f32_f16, "im2col_f32_f16", im2col_f32_f16_len, im2col_f32_f16_data, "main", 2, sizeof(vk_op_im2col_push_constants), {512, 1, 1}, { device->subgroup_size }, 1, true); -+ } -+ -+ ggml_vk_create_pipeline(device, device->pipeline_timestep_embedding_f32, "timestep_embedding_f32", timestep_embedding_f32_len, timestep_embedding_f32_data, "main", 2, sizeof(vk_op_timestep_embedding_push_constants), {256, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_pool2d_f32, "pool2d_f32", pool2d_f32_len, pool2d_f32_data, "main", 2, sizeof(vk_op_pool2d_push_constants), {512, 1, 1}, {}, 1); -+ -+ ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv6_f32, "rwkv_wkv6_f32", rwkv_wkv6_f32_len, rwkv_wkv6_f32_data, "main", 7, sizeof(vk_op_rwkv_wkv6_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); -+ -+ for (auto &c : compiles) { -+ c.wait(); -+ } -+ std::cerr << "Done!" << std::endl; -+} -+ -+static bool ggml_vk_khr_cooperative_matrix_support(const vk::PhysicalDeviceProperties& props, const vk::PhysicalDeviceDriverProperties& driver_props); -+ -+static vk_device ggml_vk_get_device(size_t idx) { -+ VK_LOG_DEBUG("ggml_vk_get_device(" << idx << ")"); -+ -+ if (vk_instance.devices[idx] == nullptr) { -+ VK_LOG_DEBUG("Initializing new vk_device"); -+ vk_device device = std::make_shared(); -+ vk_instance.devices[idx] = device; -+ -+#ifdef GGML_VULKAN_MEMORY_DEBUG -+ device->memory_logger = std::unique_ptr(new vk_memory_logger()); -+#endif -+#ifdef GGML_VULKAN_PERF -+ device->perf_logger = std::unique_ptr(new vk_perf_logger()); -+#endif -+ -+ size_t dev_num = vk_instance.device_indices[idx]; -+ -+ std::vector physical_devices = vk_instance.instance.enumeratePhysicalDevices(); -+ -+ if (dev_num >= physical_devices.size()) { -+ std::cerr << "ggml_vulkan: Device with index " << dev_num << " does not exist." << std::endl; -+ throw std::runtime_error("Device not found"); -+ } -+ -+ device->physical_device = physical_devices[dev_num]; -+ const std::vector ext_props = device->physical_device.enumerateDeviceExtensionProperties(); -+ -+ bool fp16_storage = false; -+ bool fp16_compute = false; -+ bool maintenance4_support = false; -+ bool sm_builtins = false; -+ bool amd_shader_core_properties2 = false; -+ bool pipeline_robustness = false; -+ bool coopmat2_support = false; -+ device->coopmat_support = false; -+ -+ // Check if maintenance4 is supported -+ for (const auto& properties : ext_props) { -+ if (strcmp("VK_KHR_maintenance4", properties.extensionName) == 0) { -+ maintenance4_support = true; -+ } else if (strcmp("VK_KHR_16bit_storage", properties.extensionName) == 0) { -+ fp16_storage = true; -+ } else if (strcmp("VK_KHR_shader_float16_int8", properties.extensionName) == 0) { -+ fp16_compute = true; -+ } else if (strcmp("VK_NV_shader_sm_builtins", properties.extensionName) == 0) { -+ sm_builtins = true; -+ } else if (strcmp("VK_AMD_shader_core_properties2", properties.extensionName) == 0) { -+ amd_shader_core_properties2 = true; -+ } else if (strcmp("VK_EXT_pipeline_robustness", properties.extensionName) == 0) { -+ pipeline_robustness = true; -+ } else if (strcmp("VK_EXT_subgroup_size_control", properties.extensionName) == 0) { -+ device->subgroup_size_control = true; -+ } else if (strcmp("VK_KHR_cooperative_matrix", properties.extensionName) == 0 && -+ !getenv("GGML_VK_DISABLE_COOPMAT")) { -+ device->coopmat_support = true; -+ device->coopmat_m = 0; -+ device->coopmat_n = 0; -+ device->coopmat_k = 0; -+ } else if (strcmp("VK_NV_cooperative_matrix2", properties.extensionName) == 0 && -+ !getenv("GGML_VK_DISABLE_COOPMAT2")) { -+ coopmat2_support = true; -+ } -+ } -+ -+ vk::PhysicalDeviceProperties2 props2; -+ vk::PhysicalDeviceMaintenance3Properties props3; -+ vk::PhysicalDeviceMaintenance4Properties props4; -+ vk::PhysicalDeviceSubgroupProperties subgroup_props; -+ vk::PhysicalDeviceDriverProperties driver_props; -+ vk::PhysicalDeviceShaderSMBuiltinsPropertiesNV sm_props; -+ vk::PhysicalDeviceShaderCoreProperties2AMD amd_shader_core_properties2_props; -+ vk::PhysicalDeviceVulkan12Properties vk12_props; -+ vk::PhysicalDeviceSubgroupSizeControlPropertiesEXT subgroup_size_control_props; -+ -+ props2.pNext = &props3; -+ props3.pNext = &subgroup_props; -+ subgroup_props.pNext = &driver_props; -+ driver_props.pNext = &vk12_props; -+ -+ VkBaseOutStructure * last_struct = (VkBaseOutStructure *)&vk12_props; -+ -+ if (maintenance4_support) { -+ last_struct->pNext = (VkBaseOutStructure *)&props4; -+ last_struct = (VkBaseOutStructure *)&props4; -+ } -+ if (sm_builtins) { -+ last_struct->pNext = (VkBaseOutStructure *)&sm_props; -+ last_struct = (VkBaseOutStructure *)&sm_props; -+ } -+ if (amd_shader_core_properties2) { -+ last_struct->pNext = (VkBaseOutStructure *)&amd_shader_core_properties2_props; -+ last_struct = (VkBaseOutStructure *)&amd_shader_core_properties2_props; -+ } -+ if (device->subgroup_size_control) { -+ last_struct->pNext = (VkBaseOutStructure *)&subgroup_size_control_props; -+ last_struct = (VkBaseOutStructure *)&subgroup_size_control_props; -+ } -+ -+#if defined(VK_NV_cooperative_matrix2) -+ vk::PhysicalDeviceCooperativeMatrix2PropertiesNV coopmat2_props; -+ if (coopmat2_support) { -+ last_struct->pNext = (VkBaseOutStructure *)&coopmat2_props; -+ last_struct = (VkBaseOutStructure *)&coopmat2_props; -+ } -+#endif -+ -+ device->physical_device.getProperties2(&props2); -+ device->properties = props2.properties; -+ -+ const char* GGML_VK_FORCE_MAX_ALLOCATION_SIZE = getenv("GGML_VK_FORCE_MAX_ALLOCATION_SIZE"); -+ -+ if (GGML_VK_FORCE_MAX_ALLOCATION_SIZE != nullptr) { -+ device->max_memory_allocation_size = std::stoul(GGML_VK_FORCE_MAX_ALLOCATION_SIZE); -+ } else if (maintenance4_support) { -+ device->max_memory_allocation_size = std::min(props3.maxMemoryAllocationSize, props4.maxBufferSize); -+ } else { -+ device->max_memory_allocation_size = props3.maxMemoryAllocationSize; -+ } -+ -+ device->vendor_id = device->properties.vendorID; -+ device->subgroup_size = subgroup_props.subgroupSize; -+ device->uma = device->properties.deviceType == vk::PhysicalDeviceType::eIntegratedGpu; -+ if (sm_builtins) { -+ device->shader_core_count = sm_props.shaderSMCount; -+ } else if (amd_shader_core_properties2) { -+ device->shader_core_count = amd_shader_core_properties2_props.activeComputeUnitCount; -+ } else { -+ device->shader_core_count = 0; -+ } -+ device->float_controls_rte_fp16 = vk12_props.shaderRoundingModeRTEFloat16; -+ -+ const bool force_disable_f16 = getenv("GGML_VK_DISABLE_F16") != nullptr; -+ -+ device->fp16 = !force_disable_f16 && fp16_storage && fp16_compute; -+ -+ if (!ggml_vk_khr_cooperative_matrix_support(device->properties, driver_props)) { -+ device->coopmat_support = false; -+ } -+ -+ std::vector queue_family_props = device->physical_device.getQueueFamilyProperties(); -+ -+ // Try to find a non-graphics compute queue and transfer-focused queues -+ const uint32_t compute_queue_family_index = ggml_vk_find_queue_family_index(queue_family_props, vk::QueueFlagBits::eCompute, vk::QueueFlagBits::eGraphics, -1, 1); -+ const uint32_t transfer_queue_family_index = ggml_vk_find_queue_family_index(queue_family_props, vk::QueueFlagBits::eTransfer, vk::QueueFlagBits::eCompute | vk::QueueFlagBits::eGraphics, compute_queue_family_index, 1); -+ -+ const float priorities[] = { 1.0f, 1.0f }; -+ device->single_queue = compute_queue_family_index == transfer_queue_family_index && queue_family_props[compute_queue_family_index].queueCount == 1; -+ -+ std::vector device_queue_create_infos; -+ if (compute_queue_family_index != transfer_queue_family_index) { -+ device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities}); -+ device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), transfer_queue_family_index, 1, priorities + 1}); -+ } else if(!device->single_queue) { -+ device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 2, priorities}); -+ } else { -+ device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities}); -+ } -+ vk::DeviceCreateInfo device_create_info; -+ std::vector device_extensions; -+ vk::PhysicalDeviceFeatures device_features = device->physical_device.getFeatures(); -+ -+ VkPhysicalDeviceFeatures2 device_features2; -+ device_features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2; -+ device_features2.pNext = nullptr; -+ device_features2.features = (VkPhysicalDeviceFeatures)device_features; -+ -+ VkPhysicalDeviceVulkan11Features vk11_features; -+ vk11_features.pNext = nullptr; -+ vk11_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES; -+ device_features2.pNext = &vk11_features; -+ -+ VkPhysicalDeviceVulkan12Features vk12_features; -+ vk12_features.pNext = nullptr; -+ vk12_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES; -+ vk11_features.pNext = &vk12_features; -+ -+ last_struct = (VkBaseOutStructure *)&vk12_features; -+ -+ VkPhysicalDevicePipelineRobustnessFeaturesEXT pl_robustness_features; -+ pl_robustness_features.pNext = nullptr; -+ pl_robustness_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PIPELINE_ROBUSTNESS_FEATURES_EXT; -+ pl_robustness_features.pipelineRobustness = VK_FALSE; -+ -+ if (pipeline_robustness) { -+ last_struct->pNext = (VkBaseOutStructure *)&pl_robustness_features; -+ last_struct = (VkBaseOutStructure *)&pl_robustness_features; -+ device_extensions.push_back("VK_EXT_pipeline_robustness"); -+ } -+ -+ VkPhysicalDeviceSubgroupSizeControlFeaturesEXT subgroup_size_control_features; -+ subgroup_size_control_features.pNext = nullptr; -+ subgroup_size_control_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES_EXT; -+ subgroup_size_control_features.computeFullSubgroups = false; -+ subgroup_size_control_features.subgroupSizeControl = false; -+ -+ if (device->subgroup_size_control) { -+ last_struct->pNext = (VkBaseOutStructure *)&subgroup_size_control_features; -+ last_struct = (VkBaseOutStructure *)&subgroup_size_control_features; -+ } -+ -+ VkPhysicalDeviceCooperativeMatrixFeaturesKHR coopmat_features; -+ coopmat_features.pNext = nullptr; -+ coopmat_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR; -+ coopmat_features.cooperativeMatrix = VK_FALSE; -+ -+ if (device->coopmat_support) { -+ last_struct->pNext = (VkBaseOutStructure *)&coopmat_features; -+ last_struct = (VkBaseOutStructure *)&coopmat_features; -+ } -+ -+#if defined(VK_NV_cooperative_matrix2) -+ VkPhysicalDeviceCooperativeMatrix2FeaturesNV coopmat2_features {}; -+ coopmat2_features.pNext = nullptr; -+ coopmat2_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_2_FEATURES_NV; -+ if (coopmat2_support) { -+ last_struct->pNext = (VkBaseOutStructure *)&coopmat2_features; -+ last_struct = (VkBaseOutStructure *)&coopmat2_features; -+ device_extensions.push_back("VK_NV_cooperative_matrix2"); -+ } -+#endif -+ -+ vkGetPhysicalDeviceFeatures2(device->physical_device, &device_features2); -+ -+ device->fp16 = device->fp16 && vk12_features.shaderFloat16; -+ -+ device->pipeline_robustness = pl_robustness_features.pipelineRobustness; -+ -+ if (device->subgroup_size_control) { -+ device->subgroup_min_size = subgroup_size_control_props.minSubgroupSize; -+ device->subgroup_max_size = subgroup_size_control_props.maxSubgroupSize; -+ } -+ -+ device->subgroup_size_control = device->subgroup_size_control && -+ (subgroup_size_control_props.requiredSubgroupSizeStages & vk::ShaderStageFlagBits::eCompute) && -+ subgroup_size_control_features.subgroupSizeControl; -+ -+ if (device->subgroup_size_control) { -+ device->subgroup_require_full_support = subgroup_size_control_features.computeFullSubgroups; -+ device_extensions.push_back("VK_EXT_subgroup_size_control"); -+ } -+ -+ device->coopmat_support = device->coopmat_support && coopmat_features.cooperativeMatrix; -+ -+ if (coopmat2_support) { -+#if defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) -+ if (coopmat2_features.cooperativeMatrixWorkgroupScope && -+ coopmat2_features.cooperativeMatrixFlexibleDimensions && -+ coopmat2_features.cooperativeMatrixReductions && -+ coopmat2_features.cooperativeMatrixConversions && -+ coopmat2_features.cooperativeMatrixPerElementOperations && -+ coopmat2_features.cooperativeMatrixTensorAddressing && -+ coopmat2_features.cooperativeMatrixBlockLoads && -+ vk12_features.bufferDeviceAddress) { -+ -+ std::vector flexible_dimensions; -+ uint32_t count = 0; -+ -+ PFN_vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV -+ _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV = -+ (PFN_vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV) -+ vk_instance.instance.getProcAddr("vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV"); -+ -+ _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV(device->physical_device, &count, nullptr); -+ -+ VkCooperativeMatrixFlexibleDimensionsPropertiesNV empty_prop {}; -+ empty_prop.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_FLEXIBLE_DIMENSIONS_PROPERTIES_NV; -+ flexible_dimensions.resize(count, empty_prop); -+ -+ _vkGetPhysicalDeviceCooperativeMatrixFlexibleDimensionsPropertiesNV(device->physical_device, &count, flexible_dimensions.data()); -+ -+ bool found_fp16_128 = false, -+ found_fp16_256 = false, -+ found_fp32_128 = false, -+ found_fp32_256 = false; -+ // need to support fp16*fp16 with fp16/fp32 accumulator, for workgroupsize 128 -+ // with 32x16x16 and 256 with 32x32x16. -+ for (auto &prop : flexible_dimensions) { -+ if (prop.saturatingAccumulation == VK_FALSE && -+ prop.scope == VK_SCOPE_WORKGROUP_KHR && -+ prop.AType == VK_COMPONENT_TYPE_FLOAT16_KHR && -+ prop.BType == VK_COMPONENT_TYPE_FLOAT16_KHR) { -+ -+ if (prop.workgroupInvocations == 128 && -+ prop.MGranularity <= 32 && -+ prop.NGranularity <= 16 && -+ prop.KGranularity <= 16) { -+ if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR && -+ prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) { -+ found_fp16_128 = true; -+ } -+ if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && -+ prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) { -+ found_fp32_128 = true; -+ } -+ } -+ if (prop.workgroupInvocations == 256 && -+ prop.MGranularity <= 32 && -+ prop.NGranularity <= 32 && -+ prop.KGranularity <= 16) { -+ if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR && -+ prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) { -+ found_fp16_256 = true; -+ } -+ if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && -+ prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) { -+ found_fp32_256 = true; -+ } -+ } -+ } -+ } -+ if (found_fp16_128 && found_fp16_256 && -+ found_fp32_128 && found_fp32_256 && -+ coopmat2_props.cooperativeMatrixFlexibleDimensionsMaxDimension >= 512) { -+ device->coopmat2 = true; -+ } -+ } -+#endif -+ } -+ -+ if (!vk11_features.storageBuffer16BitAccess) { -+ std::cerr << "ggml_vulkan: device " << GGML_VK_NAME << idx << " does not support 16-bit storage." << std::endl; -+ throw std::runtime_error("Unsupported device"); -+ } -+ -+ device_extensions.push_back("VK_KHR_16bit_storage"); -+ -+#ifdef GGML_VULKAN_VALIDATE -+ device_extensions.push_back("VK_KHR_shader_non_semantic_info"); -+#endif -+ -+ if (device->fp16) { -+ device_extensions.push_back("VK_KHR_shader_float16_int8"); -+ } -+ -+ if (device->coopmat_support) { -+ // Query supported shapes -+ std::vector cm_props; -+ -+ PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR = -+ (PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR)vkGetInstanceProcAddr(vk_instance.instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR"); -+ -+ uint32_t cm_props_num; -+ -+ pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR(device->physical_device, &cm_props_num, nullptr); -+ -+ cm_props.resize(cm_props_num); -+ -+ for (auto& prop : cm_props) { -+ prop.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR; -+ } -+ -+ pfn_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR(device->physical_device, &cm_props_num, cm_props.data()); -+ -+ VK_LOG_DEBUG("ggml_vulkan: Cooperative Matrix Shapes: " << cm_props.size()); -+ -+ for (auto& prop : cm_props) { -+ VK_LOG_DEBUG("ggml_vulkan: M: " << prop.MSize << " N: " << prop.NSize << " K: " << prop.KSize << " A: " << vk::to_string((vk::ComponentTypeKHR)prop.AType) << " B: " << vk::to_string((vk::ComponentTypeKHR)prop.BType) << " C: " << vk::to_string((vk::ComponentTypeKHR)prop.CType) << " Result: " << vk::to_string((vk::ComponentTypeKHR)prop.ResultType) << " saturatingAccumulation: " << prop.saturatingAccumulation << " scope: " << vk::to_string((vk::ScopeKHR)prop.scope)); -+ -+ if ((vk::ComponentTypeKHR)prop.AType == vk::ComponentTypeKHR::eFloat16 && -+ (vk::ComponentTypeKHR)prop.BType == vk::ComponentTypeKHR::eFloat16 && -+ (vk::ScopeKHR)prop.scope == vk::ScopeKHR::eSubgroup -+ ) { -+ if ((vk::ComponentTypeKHR)prop.CType == vk::ComponentTypeKHR::eFloat32 && -+ (vk::ComponentTypeKHR)prop.ResultType == vk::ComponentTypeKHR::eFloat32) { -+ // coopmat sizes not set yet -+ if (device->coopmat_m == 0) { -+ device->coopmat_acc_f32_support = true; -+ device->coopmat_m = prop.MSize; -+ device->coopmat_n = prop.NSize; -+ device->coopmat_k = prop.KSize; -+ } else if (device->coopmat_m == prop.MSize && device->coopmat_n == prop.NSize && device->coopmat_k == prop.KSize) { -+ // Only enable if shape is identical -+ device->coopmat_acc_f32_support = true; -+ } -+ } else if ((vk::ComponentTypeKHR)prop.CType == vk::ComponentTypeKHR::eFloat16 && -+ (vk::ComponentTypeKHR)prop.ResultType == vk::ComponentTypeKHR::eFloat16) { -+ // coopmat sizes not set yet -+ if (device->coopmat_m == 0) { -+ device->coopmat_acc_f16_support = true; -+ device->coopmat_m = prop.MSize; -+ device->coopmat_n = prop.NSize; -+ device->coopmat_k = prop.KSize; -+ } else if (device->coopmat_m == prop.MSize && device->coopmat_n == prop.NSize && device->coopmat_k == prop.KSize) { -+ // Only enable if shape is identical -+ device->coopmat_acc_f16_support = true; -+ } -+ } -+ } -+ } -+ -+ if (device->coopmat_m == 0 || !device->coopmat_acc_f32_support) { -+ // No suitable matmul mode found -+ GGML_LOG_DEBUG("ggml_vulkan: WARNING: No suitable matrix core mode found. Disabling matrix cores.\n"); -+ device->coopmat_support = false; -+ } -+ } -+ -+ if (device->coopmat_support) { -+ device_extensions.push_back("VK_KHR_cooperative_matrix"); -+ } -+ -+ device->name = GGML_VK_NAME + std::to_string(idx); -+ -+ device_create_info = { -+ vk::DeviceCreateFlags(), -+ device_queue_create_infos, -+ {}, -+ device_extensions -+ }; -+ device_create_info.setPNext(&device_features2); -+ device->device = device->physical_device.createDevice(device_create_info); -+ -+ // Queues -+ ggml_vk_create_queue(device, device->compute_queue, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false); -+ -+ // Shaders -+ // Disable matmul tile sizes early if performance low or not supported -+ switch (device->vendor_id) { -+#ifndef GGML_VULKAN_RUN_TESTS -+ case VK_VENDOR_ID_AMD: -+ case VK_VENDOR_ID_INTEL: -+ device->mul_mat_l = false; -+ device->mul_mat_m = true; -+ device->mul_mat_s = true; -+ device->mul_mat_id_l = false; -+ device->mul_mat_id_m = true; -+ device->mul_mat_id_s = true; -+ break; -+ case VK_VENDOR_ID_APPLE: -+ device->mul_mat_l = false; -+ device->mul_mat_m = true; -+ device->mul_mat_s = false; -+ device->mul_mat_id_l = false; -+ device->mul_mat_id_m = true; -+ device->mul_mat_id_s = false; -+ break; -+#endif -+ default: -+ device->mul_mat_l = true; -+ device->mul_mat_m = true; -+ device->mul_mat_s = true; -+ device->mul_mat_id_l = true; -+ device->mul_mat_id_m = true; -+ device->mul_mat_id_s = true; -+ break; -+ } -+ -+ ggml_vk_load_shaders(device); -+ -+ if (!device->single_queue) { -+ const uint32_t transfer_queue_index = compute_queue_family_index == transfer_queue_family_index ? 1 : 0; -+ ggml_vk_create_queue(device, device->transfer_queue, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true); -+ } else { -+ // TODO: Use pointer or reference to avoid copy -+ device->transfer_queue = device->compute_queue; -+ } -+ -+ device->buffer_type = { -+ /* .iface = */ ggml_backend_vk_buffer_type_interface, -+ /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), idx), -+ /* .context = */ new ggml_backend_vk_buffer_type_context{ device->name, device }, -+ }; -+ -+ device->fence = device->device.createFence({}); -+ -+ device->idx = idx; -+ -+ return device; -+ } -+ -+ return vk_instance.devices[idx]; -+} -+ -+static void ggml_vk_print_gpu_info(size_t idx) { -+ GGML_ASSERT(idx < vk_instance.device_indices.size()); -+ size_t dev_num = vk_instance.device_indices[idx]; -+ VK_LOG_DEBUG("ggml_vk_print_gpu_info(" << dev_num << ")"); -+ GGML_ASSERT(vk_instance_initialized); -+ -+ std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); -+ -+ if (dev_num >= devices.size()) { -+ std::cerr << "ggml_vulkan: Device with index " << dev_num << " does not exist." << std::endl; -+ throw std::runtime_error("Device not found"); -+ } -+ -+ vk::PhysicalDevice physical_device = devices[dev_num]; -+ std::vector ext_props = physical_device.enumerateDeviceExtensionProperties(); -+ -+ vk::PhysicalDeviceProperties2 props2; -+ vk::PhysicalDeviceMaintenance3Properties props3; -+ vk::PhysicalDeviceSubgroupProperties subgroup_props; -+ vk::PhysicalDeviceDriverProperties driver_props; -+ props2.pNext = &props3; -+ props3.pNext = &subgroup_props; -+ subgroup_props.pNext = &driver_props; -+ physical_device.getProperties2(&props2); -+ -+ const size_t subgroup_size = subgroup_props.subgroupSize; -+ const bool uma = props2.properties.deviceType == vk::PhysicalDeviceType::eIntegratedGpu; -+ -+ bool fp16_storage = false; -+ bool fp16_compute = false; -+ bool coopmat_support = false; -+ bool coopmat2_support = false; -+ -+ for (auto properties : ext_props) { -+ if (strcmp("VK_KHR_16bit_storage", properties.extensionName) == 0) { -+ fp16_storage = true; -+ } else if (strcmp("VK_KHR_shader_float16_int8", properties.extensionName) == 0) { -+ fp16_compute = true; -+ } else if (strcmp("VK_KHR_cooperative_matrix", properties.extensionName) == 0 && -+ !getenv("GGML_VK_DISABLE_COOPMAT")) { -+ coopmat_support = true; -+#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) -+ } else if (strcmp("VK_NV_cooperative_matrix2", properties.extensionName) == 0 && -+ !getenv("GGML_VK_DISABLE_COOPMAT2")) { -+ coopmat2_support = true; -+#endif -+ } -+ } -+ -+ if (!ggml_vk_khr_cooperative_matrix_support(props2.properties, driver_props)) { -+ coopmat_support = false; -+ } -+ -+ const char* GGML_VK_DISABLE_F16 = getenv("GGML_VK_DISABLE_F16"); -+ bool force_disable_f16 = GGML_VK_DISABLE_F16 != nullptr; -+ -+ bool fp16 = !force_disable_f16 && fp16_storage && fp16_compute; -+ -+ vk::PhysicalDeviceFeatures device_features = physical_device.getFeatures(); -+ -+ VkPhysicalDeviceFeatures2 device_features2; -+ device_features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2; -+ device_features2.pNext = nullptr; -+ device_features2.features = (VkPhysicalDeviceFeatures)device_features; -+ -+ VkPhysicalDeviceVulkan11Features vk11_features; -+ vk11_features.pNext = nullptr; -+ vk11_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES; -+ device_features2.pNext = &vk11_features; -+ -+ VkPhysicalDeviceVulkan12Features vk12_features; -+ vk12_features.pNext = nullptr; -+ vk12_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES; -+ vk11_features.pNext = &vk12_features; -+ -+ // Pointer to the last chain element -+ VkBaseOutStructure * last_struct = (VkBaseOutStructure *)&vk12_features; -+ -+ VkPhysicalDeviceCooperativeMatrixFeaturesKHR coopmat_features; -+ coopmat_features.pNext = nullptr; -+ coopmat_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR; -+ coopmat_features.cooperativeMatrix = VK_FALSE; -+ -+ if (coopmat_support) { -+ last_struct->pNext = (VkBaseOutStructure *)&coopmat_features; -+ last_struct = (VkBaseOutStructure *)&coopmat_features; -+ } -+ -+ vkGetPhysicalDeviceFeatures2(physical_device, &device_features2); -+ -+ fp16 = fp16 && vk12_features.shaderFloat16; -+ -+ coopmat_support = coopmat_support && coopmat_features.cooperativeMatrix; -+ -+ std::string matrix_cores = coopmat2_support ? "NV_coopmat2" : coopmat_support ? "KHR_coopmat" : "none"; -+ -+ std::string device_name = props2.properties.deviceName.data(); -+ GGML_LOG_DEBUG("ggml_vulkan: %zu = %s (%s) | uma: %d | fp16: %d | warp size: %zu | matrix cores: %s\n", -+ idx, device_name.c_str(), driver_props.driverName.data(), uma, fp16, subgroup_size, matrix_cores.c_str()); -+ -+ if (props2.properties.deviceType == vk::PhysicalDeviceType::eCpu) { -+ GGML_LOG_DEBUG("ggml_vulkan: Warning: Device type is CPU. This is probably not the device you want.\n"); -+ } -+} -+ -+static bool ggml_vk_instance_validation_ext_available(const std::vector& instance_extensions); -+static bool ggml_vk_instance_portability_enumeration_ext_available(const std::vector& instance_extensions); -+ -+void ggml_vk_instance_init() { -+ if (vk_instance_initialized) { -+ return; -+ } -+ VK_LOG_DEBUG("ggml_vk_instance_init()"); -+ -+ vk_instance_initialized = true; -+ -+ vk::ApplicationInfo app_info{ "ggml-vulkan", 1, nullptr, 0, VK_API_VERSION }; -+ -+ const std::vector instance_extensions = vk::enumerateInstanceExtensionProperties(); -+ const bool validation_ext = ggml_vk_instance_validation_ext_available(instance_extensions); -+#ifdef __APPLE__ -+ const bool portability_enumeration_ext = ggml_vk_instance_portability_enumeration_ext_available(instance_extensions); -+#endif -+ -+ std::vector layers; -+ -+ if (validation_ext) { -+ layers.push_back("VK_LAYER_KHRONOS_validation"); -+ } -+ std::vector extensions; -+ if (validation_ext) { -+ extensions.push_back("VK_EXT_validation_features"); -+ } -+#ifdef __APPLE__ -+ if (portability_enumeration_ext) { -+ extensions.push_back("VK_KHR_portability_enumeration"); -+ } -+#endif -+ vk::InstanceCreateInfo instance_create_info(vk::InstanceCreateFlags{}, &app_info, layers, extensions); -+#ifdef __APPLE__ -+ if (portability_enumeration_ext) { -+ instance_create_info.flags |= vk::InstanceCreateFlagBits::eEnumeratePortabilityKHR; -+ } -+#endif -+ -+ std::vector features_enable; -+ vk::ValidationFeaturesEXT validation_features; -+ -+ if (validation_ext) { -+ features_enable = { vk::ValidationFeatureEnableEXT::eBestPractices }; -+ validation_features = { -+ features_enable, -+ {}, -+ }; -+ validation_features.setPNext(nullptr); -+ instance_create_info.setPNext(&validation_features); -+ GGML_LOG_DEBUG("ggml_vulkan: Validation layers enabled\n"); -+ } -+ vk_instance.instance = vk::createInstance(instance_create_info); -+ -+ size_t num_available_devices = vk_instance.instance.enumeratePhysicalDevices().size(); -+ -+ // Emulate behavior of CUDA_VISIBLE_DEVICES for Vulkan -+ char * devices_env = getenv("GGML_VK_VISIBLE_DEVICES"); -+ if (devices_env != nullptr) { -+ std::string devices(devices_env); -+ std::replace(devices.begin(), devices.end(), ',', ' '); -+ -+ std::stringstream ss(devices); -+ size_t tmp; -+ while (ss >> tmp) { -+ if(tmp >= num_available_devices) { -+ std::cerr << "ggml_vulkan: Invalid device index " << tmp << " in GGML_VK_VISIBLE_DEVICES." << std::endl; -+ throw std::runtime_error("Invalid Vulkan device index"); -+ } -+ vk_instance.device_indices.push_back(tmp); -+ } -+ } else { -+ std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); -+ -+ // Make sure at least one device exists -+ if (devices.empty()) { -+ std::cerr << "ggml_vulkan: Error: No devices found." << std::endl; -+ GGML_ABORT("fatal error"); -+ } -+ -+ // Default to using all dedicated GPUs -+ for (size_t i = 0; i < devices.size(); i++) { -+ vk::PhysicalDeviceProperties2 new_props; -+ vk::PhysicalDeviceDriverProperties new_driver; -+ vk::PhysicalDeviceIDProperties new_id; -+ new_props.pNext = &new_driver; -+ new_driver.pNext = &new_id; -+ devices[i].getProperties2(&new_props); -+ -+ if (new_props.properties.deviceType == vk::PhysicalDeviceType::eDiscreteGpu) { -+ // Check if there are two physical devices corresponding to the same GPU -+ auto old_device = std::find_if( -+ vk_instance.device_indices.begin(), -+ vk_instance.device_indices.end(), -+ [&devices, &new_id](const size_t k){ -+ vk::PhysicalDeviceProperties2 old_props; -+ vk::PhysicalDeviceIDProperties old_id; -+ old_props.pNext = &old_id; -+ devices[k].getProperties2(&old_props); -+ return std::equal(std::begin(old_id.deviceUUID), std::end(old_id.deviceUUID), std::begin(new_id.deviceUUID)); -+ } -+ ); -+ if (old_device == vk_instance.device_indices.end()) { -+ vk_instance.device_indices.push_back(i); -+ } else { -+ // There can be two physical devices corresponding to the same GPU if there are 2 different drivers -+ // This can cause error when splitting layers aross the devices, need to keep only 1 -+ VK_LOG_DEBUG("Device " << i << " and device " << *old_device << " have the same deviceUUID"); -+ -+ vk::PhysicalDeviceProperties2 old_props; -+ vk::PhysicalDeviceDriverProperties old_driver; -+ old_props.pNext = &old_driver; -+ devices[*old_device].getProperties2(&old_props); -+ -+ std::map driver_priorities {}; -+ int old_priority = std::numeric_limits::max(); -+ int new_priority = std::numeric_limits::max(); -+ -+ // Check https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/VkDriverId.html for the list of driver id -+ // Smaller number -> higher priority -+ switch (old_props.properties.vendorID) { -+ case VK_VENDOR_ID_AMD: -+ driver_priorities[vk::DriverId::eMesaRadv] = 1; -+ driver_priorities[vk::DriverId::eAmdOpenSource] = 2; -+ driver_priorities[vk::DriverId::eAmdProprietary] = 3; -+ break; -+ case VK_VENDOR_ID_INTEL: -+ driver_priorities[vk::DriverId::eIntelOpenSourceMESA] = 1; -+ driver_priorities[vk::DriverId::eIntelProprietaryWindows] = 2; -+ break; -+ case VK_VENDOR_ID_NVIDIA: -+ driver_priorities[vk::DriverId::eNvidiaProprietary] = 1; -+#if defined(VK_API_VERSION_1_3) && VK_HEADER_VERSION >= 235 -+ driver_priorities[vk::DriverId::eMesaNvk] = 2; -+#endif -+ break; -+ } -+ -+ if (driver_priorities.count(old_driver.driverID)) { -+ old_priority = driver_priorities[old_driver.driverID]; -+ } -+ if (driver_priorities.count(new_driver.driverID)) { -+ new_priority = driver_priorities[new_driver.driverID]; -+ } -+ -+ if (new_priority < old_priority) { -+ auto r = std::remove(vk_instance.device_indices.begin(), vk_instance.device_indices.end(), *old_device); -+ vk_instance.device_indices.erase(r, vk_instance.device_indices.end()); -+ vk_instance.device_indices.push_back(i); -+ -+ VK_LOG_DEBUG("Prioritize device " << i << " driver " << new_driver.driverName << " over device " << *old_device << " driver " << old_driver.driverName); -+ } -+ else { -+ VK_LOG_DEBUG("Prioritize device " << *old_device << " driver " << old_driver.driverName << " over device " << i << " driver " << new_driver.driverName << std::endl); -+ } -+ } -+ } -+ } -+ -+ // If no dedicated GPUs found, fall back to GPU 0 -+ if (vk_instance.device_indices.empty()) { -+ vk_instance.device_indices.push_back(0); -+ } -+ } -+ GGML_LOG_DEBUG("ggml_vulkan: Found %zu Vulkan devices:\n", vk_instance.device_indices.size()); -+ -+ for (size_t i = 0; i < vk_instance.device_indices.size(); i++) { -+ ggml_vk_print_gpu_info(i); -+ } -+} -+ -+static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) { -+ VK_LOG_DEBUG("ggml_vk_init(" << ctx->name << ", " << idx << ")"); -+ ggml_vk_instance_init(); -+ GGML_ASSERT(idx < vk_instance.device_indices.size()); -+ -+ ctx->name = GGML_VK_NAME + std::to_string(idx); -+ -+ ctx->device = ggml_vk_get_device(idx); -+ -+ ctx->semaphore_idx = 0; -+ ctx->event_idx = 0; -+ -+ ctx->prealloc_size_x = 0; -+ ctx->prealloc_size_y = 0; -+ ctx->prealloc_size_split_k = 0; -+ -+ ctx->fence = ctx->device->device.createFence({}); -+ -+#ifdef GGML_VULKAN_CHECK_RESULTS -+ const char* skip_checks = getenv("GGML_VULKAN_SKIP_CHECKS"); -+ vk_skip_checks = (skip_checks == NULL ? 0 : atoi(skip_checks)); -+ const char* output_tensor = getenv("GGML_VULKAN_OUTPUT_TENSOR"); -+ vk_output_tensor = (output_tensor == NULL ? 0 : atoi(output_tensor)); -+#endif -+} -+ -+static vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type type) { -+ VK_LOG_DEBUG("ggml_vk_get_to_fp16()"); -+ switch (type) { -+ case GGML_TYPE_F32: -+ case GGML_TYPE_Q4_0: -+ case GGML_TYPE_Q4_1: -+ case GGML_TYPE_Q5_0: -+ case GGML_TYPE_Q5_1: -+ case GGML_TYPE_Q8_0: -+ case GGML_TYPE_Q2_K: -+ case GGML_TYPE_Q3_K: -+ case GGML_TYPE_Q4_K: -+ case GGML_TYPE_Q5_K: -+ case GGML_TYPE_Q6_K: -+ case GGML_TYPE_IQ4_NL: -+ break; -+ default: -+ return nullptr; -+ } -+ -+ return ctx->device->pipeline_dequant[type]; -+} -+ -+static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec) { -+ VK_LOG_DEBUG("ggml_vk_get_mul_mat_mat_pipeline(" << ggml_type_name(src0_type) << ", " << ggml_type_name(src1_type) << ")"); -+ if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_matmul_f32; -+ } -+ if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F16) { -+ return ctx->device->pipeline_matmul_f32_f16; -+ } -+ if (prec == GGML_PREC_DEFAULT && ctx->device->fp16 && !(ctx->device->coopmat_support && !ctx->device->coopmat_acc_f16_support)) { -+ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_matmul_f16_f32.f16acc; -+ } -+ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { -+ return ctx->device->pipeline_matmul_f16.f16acc; -+ } -+ } else { -+ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_matmul_f16_f32.f32acc; -+ } -+ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { -+ return ctx->device->pipeline_matmul_f16.f32acc; -+ } -+ } -+ -+ if (src1_type != GGML_TYPE_F32 && !ctx->device->coopmat2) { -+ return nullptr; -+ } -+ -+ switch (src0_type) { -+ case GGML_TYPE_Q4_0: -+ case GGML_TYPE_Q4_1: -+ case GGML_TYPE_Q5_0: -+ case GGML_TYPE_Q5_1: -+ case GGML_TYPE_Q8_0: -+ case GGML_TYPE_Q2_K: -+ case GGML_TYPE_Q3_K: -+ case GGML_TYPE_Q4_K: -+ case GGML_TYPE_Q5_K: -+ case GGML_TYPE_Q6_K: -+ case GGML_TYPE_IQ4_NL: -+ break; -+ default: -+ return nullptr; -+ } -+ -+ if (ctx->device->coopmat2) { -+ assert(src1_type == GGML_TYPE_F16); -+ return ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f16acc; -+ } -+ return ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f32acc; -+} -+ -+static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * ctx, ggml_type a_type, ggml_type b_type, uint32_t num_cols) { -+ VK_LOG_DEBUG("ggml_vk_get_dequantize_mul_mat_vec()"); -+ GGML_ASSERT(b_type == GGML_TYPE_F32 || b_type == GGML_TYPE_F16); -+ GGML_ASSERT(num_cols >= 1 && num_cols <= mul_mat_vec_max_cols); -+ -+ switch (a_type) { -+ case GGML_TYPE_F32: -+ case GGML_TYPE_F16: -+ case GGML_TYPE_Q4_0: -+ case GGML_TYPE_Q4_1: -+ case GGML_TYPE_Q5_0: -+ case GGML_TYPE_Q5_1: -+ case GGML_TYPE_Q8_0: -+ case GGML_TYPE_Q2_K: -+ case GGML_TYPE_Q3_K: -+ case GGML_TYPE_Q4_K: -+ case GGML_TYPE_Q5_K: -+ case GGML_TYPE_Q6_K: -+ case GGML_TYPE_IQ4_NL: -+ break; -+ default: -+ return nullptr; -+ } -+ -+ return b_type == GGML_TYPE_F32 ? ctx->device->pipeline_dequant_mul_mat_vec_f32_f32[a_type][num_cols-1] : ctx->device->pipeline_dequant_mul_mat_vec_f16_f32[a_type][num_cols-1]; -+} -+ -+static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec) { -+ VK_LOG_DEBUG("ggml_vk_get_mul_mat_mat_id_pipeline()"); -+ if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_matmul_id_f32; -+ } -+ if (prec == GGML_PREC_DEFAULT && ctx->device->fp16 && !(ctx->device->coopmat_support && !ctx->device->coopmat_acc_f16_support)) { -+ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_matmul_id_f16_f32.f16acc; -+ } -+ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { -+ return ctx->device->pipeline_matmul_id_f16.f16acc; -+ } -+ } else { -+ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_matmul_id_f16_f32.f32acc; -+ } -+ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { -+ return ctx->device->pipeline_matmul_id_f16.f32acc; -+ } -+ } -+ -+ GGML_ASSERT(src1_type == GGML_TYPE_F32); -+ -+ switch (src0_type) { -+ case GGML_TYPE_Q4_0: -+ case GGML_TYPE_Q4_1: -+ case GGML_TYPE_Q5_0: -+ case GGML_TYPE_Q5_1: -+ case GGML_TYPE_Q8_0: -+ case GGML_TYPE_Q2_K: -+ case GGML_TYPE_Q3_K: -+ case GGML_TYPE_Q4_K: -+ case GGML_TYPE_Q5_K: -+ case GGML_TYPE_Q6_K: -+ case GGML_TYPE_IQ4_NL: -+ break; -+ default: -+ return nullptr; -+ } -+ -+ return ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat_id[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat_id[src0_type].f32acc; -+} -+ -+static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context * ctx, ggml_type a_type, ggml_type b_type) { -+ VK_LOG_DEBUG("ggml_vk_get_dequantize_mul_mat_vec()"); -+ GGML_ASSERT(b_type == GGML_TYPE_F32); -+ -+ switch (a_type) { -+ case GGML_TYPE_F32: -+ case GGML_TYPE_F16: -+ case GGML_TYPE_Q4_0: -+ case GGML_TYPE_Q4_1: -+ case GGML_TYPE_Q5_0: -+ case GGML_TYPE_Q5_1: -+ case GGML_TYPE_Q8_0: -+ case GGML_TYPE_Q2_K: -+ case GGML_TYPE_Q3_K: -+ case GGML_TYPE_Q4_K: -+ case GGML_TYPE_Q5_K: -+ case GGML_TYPE_Q6_K: -+ case GGML_TYPE_IQ4_NL: -+ break; -+ default: -+ return nullptr; -+ } -+ -+ return ctx->device->pipeline_dequant_mul_mat_vec_id_f32[a_type]; -+} -+ -+static vk_buffer ggml_vk_pool_malloc(ggml_backend_vk_context * ctx, size_t size) { -+ VK_LOG_DEBUG("ggml_vk_pool_malloc(" << size << ")"); -+ VK_LOG_MEMORY("ggml_vk_pool_malloc"); -+ -+ int best_i = -1; -+ size_t best_size = std::numeric_limits::max(); //smallest unused buffer that fits our needs -+ int worst_i = -1; -+ size_t worst_size = 0; //largest unused buffer seen so far -+ for (int i = 0; i < MAX_VK_BUFFERS; ++i) { -+ vk_buffer &b = ctx->buffer_pool[i]; -+ if (b != nullptr && b->size >= size && b->size < best_size) { -+ best_i = i; -+ best_size = b->size; -+ } -+ if (b != nullptr && b->size > worst_size) { -+ worst_i = i; -+ worst_size = b->size; -+ } -+ } -+ if(best_i != -1) { -+ //found the smallest buffer that fits our needs -+ vk_buffer b = ctx->buffer_pool[best_i]; -+ ctx->buffer_pool[best_i].reset(); -+ return b; -+ } -+ if(worst_i != -1) { -+ //no buffer that fits our needs, resize largest one to save memory -+ vk_buffer& b = ctx->buffer_pool[worst_i]; -+ ggml_vk_destroy_buffer(b); -+ } -+ -+ return ggml_vk_create_buffer_device(ctx->device, size); -+} -+ -+static void ggml_vk_pool_free(ggml_backend_vk_context * ctx, vk_buffer& buffer) { -+ VK_LOG_DEBUG("ggml_vk_pool_free(" << buffer->size << ")"); -+ for (int i = 0; i < MAX_VK_BUFFERS; ++i) { -+ vk_buffer& b = ctx->buffer_pool[i]; -+ if (b == nullptr) { -+ b = buffer; -+ return; -+ } -+ } -+ std::cerr << "ggml_vulkan: WARNING: vk buffer pool full, increase MAX_VK_BUFFERS" << std::endl; -+ ggml_vk_destroy_buffer(buffer); -+} -+ -+// Returns an available temporary buffer that may only be used temporarily, it will be reused -+static vk_buffer ggml_vk_create_buffer_temp(ggml_backend_vk_context * ctx, size_t size) { -+ // Try to find existing temp buffer with enough capacity -+ for (auto& buffer : ctx->gc.temp_buffers) { -+ if (buffer->size >= size) { -+ return buffer; -+ } -+ } -+ -+ VK_LOG_MEMORY("ggml_vk_create_buffer_temp(" << size << ")"); -+ -+ // Otherwise create new buffer -+ vk_buffer buf = ggml_vk_pool_malloc(ctx, size); -+ ctx->gc.temp_buffers.push_back(buf); -+ -+ return buf; -+} -+ -+static void * ggml_vk_host_malloc(vk_device& device, size_t size) { -+ VK_LOG_MEMORY("ggml_vk_host_malloc(" << size << ")"); -+ vk_buffer buf = ggml_vk_create_buffer(device, size, -+ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached, -+ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); -+ -+ if(!(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible)) { -+ fprintf(stderr, "WARNING: failed to allocate %.2f MB of pinned memory\n", -+ size/1024.0/1024.0); -+ device->device.freeMemory(buf->device_memory); -+ device->device.destroyBuffer(buf->buffer); -+ return nullptr; -+ } -+ -+ device->pinned_memory.push_back(std::make_tuple(buf->ptr, size, buf)); -+ -+ return buf->ptr; -+} -+ -+static void ggml_vk_host_free(vk_device& device, void* ptr) { -+ if (ptr == nullptr) { -+ return; -+ } -+ VK_LOG_MEMORY("ggml_vk_host_free(" << ptr << ")"); -+ vk_buffer buf; -+ size_t index; -+ for (size_t i = 0; i < device->pinned_memory.size(); i++) { -+ const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]); -+ const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]); -+ if (ptr >= addr && ptr < endr) { -+ buf = std::get<2>(device->pinned_memory[i]); -+ index = i; -+ break; -+ } -+ } -+ if (buf == nullptr) { -+ fprintf(stderr, "WARNING: failed to free pinned memory: memory not in map\n"); -+ return; -+ } -+ -+ ggml_vk_destroy_buffer(buf); -+ -+ device->pinned_memory.erase(device->pinned_memory.begin() + index); -+} -+ -+static void ggml_vk_host_get(vk_device& device, const void * ptr, vk_buffer& buf, size_t& buf_offset) { -+ buf = nullptr; -+ buf_offset = 0; -+ for (size_t i = 0; i < device->pinned_memory.size(); i++) { -+ const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]); -+ const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]); -+ if (ptr >= addr && ptr < endr) { -+ buf = std::get<2>(device->pinned_memory[i]); -+ buf_offset = ((const uint8_t *)ptr) - addr; -+ break; -+ } -+ } -+} -+ -+static vk_submission ggml_vk_begin_submission(vk_device& device, vk_queue& q, bool one_time = true) { -+ vk_submission s; -+ s.buffer = ggml_vk_create_cmd_buffer(device, q); -+ if (one_time) { -+ s.buffer.begin({ vk::CommandBufferUsageFlagBits::eOneTimeSubmit }); -+ } else { -+ s.buffer.begin({ vk::CommandBufferUsageFlags{} }); -+ } -+ -+ return s; -+} -+ -+ -+ -+static void ggml_vk_dispatch_pipeline(ggml_backend_vk_context* ctx, vk_context& subctx, vk_pipeline& pipeline, std::initializer_list const& descriptor_buffer_infos, size_t push_constant_size, const void* push_constants, std::array elements) { -+ const uint32_t wg0 = CEIL_DIV(elements[0], pipeline->wg_denoms[0]); -+ const uint32_t wg1 = CEIL_DIV(elements[1], pipeline->wg_denoms[1]); -+ const uint32_t wg2 = CEIL_DIV(elements[2], pipeline->wg_denoms[2]); -+ VK_LOG_DEBUG("ggml_vk_dispatch_pipeline(" << pipeline->name << ", {"; -+ for (auto& buffer : descriptor_buffer_infos) { -+ std::cerr << "(" << buffer.buffer << ", " << buffer.offset << ", " << buffer.range << "), "; -+ } -+ std::cerr << "}, (" << wg0 << "," << wg1 << "," << wg2 << "))"); -+ GGML_ASSERT(pipeline->descriptor_set_idx < pipeline->descriptor_sets.size()); -+ GGML_ASSERT(descriptor_buffer_infos.size() == pipeline->parameter_count); -+ -+ vk::DescriptorSet& descriptor_set = pipeline->descriptor_sets[pipeline->descriptor_set_idx++]; -+ vk::WriteDescriptorSet write_descriptor_set{ descriptor_set, 0, 0, pipeline->parameter_count, vk::DescriptorType::eStorageBuffer, nullptr, descriptor_buffer_infos.begin() }; -+ ctx->device->device.updateDescriptorSets({ write_descriptor_set }, {}); -+ -+ subctx->s->buffer.pushConstants(pipeline->layout, vk::ShaderStageFlagBits::eCompute, 0, push_constant_size, push_constants); -+ subctx->s->buffer.bindPipeline(vk::PipelineBindPoint::eCompute, pipeline->pipeline); -+ subctx->s->buffer.bindDescriptorSets(vk::PipelineBindPoint::eCompute, -+ pipeline->layout, -+ 0, -+ { descriptor_set }, -+ {}); -+ subctx->s->buffer.dispatch(wg0, wg1, wg2); -+} -+ -+static void ggml_vk_end_submission(vk_submission& s, std::vector wait_semaphores, std::vector signal_semaphores) { -+ s.buffer.end(); -+ -+ s.wait_semaphores = std::move(wait_semaphores); -+ s.signal_semaphores = std::move(signal_semaphores); -+} -+ -+static void ggml_vk_ctx_end(vk_context& ctx) { -+ VK_LOG_DEBUG("ggml_vk_ctx_end(" << ctx << ", " << ctx->seqs.size() << ")"); -+ if (ctx->s == nullptr) { -+ return; -+ } -+ -+ ctx->s->buffer.end(); -+ ctx->s = nullptr; -+} -+ -+static void ggml_vk_ctx_begin(vk_device& device, vk_context& subctx) { -+ VK_LOG_DEBUG("ggml_vk_ctx_begin(" << device->name << ")"); -+ if (subctx->s != nullptr) { -+ ggml_vk_ctx_end(subctx); -+ } -+ -+ subctx->seqs.push_back({ ggml_vk_begin_submission(device, *subctx->q) }); -+ subctx->s = subctx->seqs[subctx->seqs.size() - 1].data(); -+} -+ -+static size_t ggml_vk_align_size(size_t width, size_t align) { -+ VK_LOG_DEBUG("ggml_vk_align_size(" << width << ", " << align << ")"); -+ return CEIL_DIV(width, align) * align; -+} -+ -+static void deferred_memcpy(void * dst, const void * src, size_t size, std::vector* memcpys = nullptr) { -+ if (memcpys == nullptr) { -+ memcpy(dst, src, size); -+ } else { -+ memcpys->emplace_back(dst, src, size); -+ } -+} -+ -+static void ggml_vk_ensure_sync_staging_buffer(vk_device& device, size_t size) { -+ if (device->sync_staging == nullptr || device->sync_staging->size < size) { -+ VK_LOG_MEMORY("ggml_vk_ensure_sync_staging_buffer(" << size << ")"); -+ ggml_vk_destroy_buffer(device->sync_staging); -+ device->sync_staging = ggml_vk_create_buffer_check(device, size, -+ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached, -+ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent); -+ } -+} -+ -+static void ggml_vk_buffer_write_nc_async(ggml_backend_vk_context * ctx, vk_context& subctx, vk_buffer& dst, size_t offset, const ggml_tensor * tensor, bool sync_staging = false) { -+ VK_LOG_DEBUG("ggml_vk_buffer_write_nc_async(" << tensor << ")"); -+ GGML_ASSERT(!ggml_is_contiguous(tensor)); -+ // Buffer is already mapped -+ if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { -+ std::cerr << "ggml_vulkan: buffer_write_nc_async dst buffer is host_visible. Use synchronous write." << std::endl; -+ GGML_ABORT("fatal error"); -+ } -+ // Check if src is pinned memory -+ vk_buffer buf = nullptr; -+ size_t buf_offset = 0; -+ ggml_vk_host_get(ctx->device, tensor->data, buf, buf_offset); -+ -+ const uint64_t ne0 = tensor->ne[0]; -+ const uint64_t ne1 = tensor->ne[1]; -+ const uint64_t ne2 = tensor->ne[2]; -+ const uint64_t ne3 = tensor->ne[3]; -+ const uint64_t nb0 = tensor->nb[0]; -+ const uint64_t nb1 = tensor->nb[1]; -+ const uint64_t nb2 = tensor->nb[2]; -+ const uint64_t nb3 = tensor->nb[3]; -+ const ggml_type type = tensor->type; -+ const uint64_t ts = ggml_type_size(type); -+ const uint64_t bs = ggml_blck_size(type); -+ -+ const uint64_t dstnb0 = ts; -+ const uint64_t dstnb1 = dstnb0*(ne0/bs); -+ const uint64_t dstnb2 = dstnb1*ne1; -+ const uint64_t dstnb3 = dstnb2*ne2; -+ -+ const uint64_t ne = ggml_nelements(tensor); -+ -+ if (buf != nullptr) { -+ // Memory is pinned, use as staging buffer -+ std::vector slices; -+ -+ for (uint64_t i3 = 0; i3 < ne3; i3++) { -+ for (uint64_t i2 = 0; i2 < ne2; i2++) { -+ // Find longest contiguous slice -+ if (ne1*nb1 == dstnb2) { -+ slices.push_back({ buf_offset + i3*nb3 + i2*nb2, offset + i3*dstnb3 + i2*dstnb2, dstnb2 }); -+ } else { -+ for (uint64_t i1 = 0; i1 < ne1; i1++) { -+ if (ne0*nb0/bs == dstnb1) { -+ slices.push_back({ buf_offset + i3*nb3 + i2*nb2 + i1*nb1, offset + i3*dstnb3 + i2*dstnb2 + i1*dstnb1, dstnb1 }); -+ } else { -+ const uint64_t s_off = buf_offset + i3*nb3 + i2*nb2 + i1*nb1; -+ const uint64_t d_off = offset + i3*dstnb3 + i2*dstnb2 + i1*dstnb1; -+ for (uint64_t i0 = 0; i0 < ne0; i0++) { -+ slices.push_back({ s_off + i1*nb0, d_off + i0*dstnb0, dstnb0 }); -+ } -+ } -+ } -+ } -+ } -+ } -+ -+ ggml_vk_sync_buffers(subctx); -+ subctx->s->buffer.copyBuffer(buf->buffer, dst->buffer, slices); -+ return; -+ } -+ -+ if (!sync_staging) { -+ GGML_ABORT("Asynchronous write to non-pinned memory not supported"); -+ } -+ -+ // Staging buffer required -+ vk_buffer& staging = ctx->device->sync_staging; -+ const uint64_t copy_size = ts*ne/bs; -+ ggml_vk_ensure_sync_staging_buffer(ctx->device, copy_size); -+ VkBufferCopy buf_copy{ 0, offset, copy_size }; -+ -+ ggml_vk_sync_buffers(subctx); -+ vkCmdCopyBuffer(subctx->s->buffer, (VkBuffer)staging->buffer, (VkBuffer)dst->buffer, 1, &buf_copy); -+ -+ for (uint64_t i3 = 0; i3 < ne3; i3++) { -+ for (uint64_t i2 = 0; i2 < ne2; i2++) { -+ // Find longest contiguous slice -+ if (ne1*nb1 == dstnb2) { -+ deferred_memcpy((uint8_t *)staging->ptr + i3*dstnb3 + i2*dstnb2, (const uint8_t *) tensor->data + buf_offset + i3*nb3 + i2*nb2, dstnb2, &subctx->in_memcpys); -+ } else { -+ for (uint64_t i1 = 0; i1 < ne1; i1++) { -+ if (ne0*nb0/bs == dstnb1) { -+ deferred_memcpy((uint8_t *)staging->ptr + i3*dstnb3 + i2*dstnb2 + i1*dstnb1, (const uint8_t *) tensor->data + buf_offset + i3*nb3 + i2*nb2 + i1*nb1, dstnb1, &subctx->in_memcpys); -+ } else { -+ const uint64_t s_off = buf_offset + i3*nb3 + i2*nb2 + i1*nb1; -+ const uint64_t d_off = i3*dstnb3 + i2*dstnb2 + i1*dstnb1; -+ for (uint64_t i0 = 0; i0 < ne0; i0++) { -+ deferred_memcpy((uint8_t *)staging->ptr + d_off + i0*dstnb0, (const uint8_t *) tensor->data + s_off + i0*nb0, dstnb0, &subctx->in_memcpys); -+ } -+ } -+ } -+ } -+ } -+ } -+} -+ -+static void ggml_vk_buffer_write_2d_async(vk_context subctx, vk_buffer& dst, size_t offset, const void * src, size_t spitch, size_t width, size_t height, bool sync_staging = false) { -+ VK_LOG_DEBUG("ggml_vk_buffer_write_2d_async(" << width << ", " << height << ")"); -+ // Buffer is already mapped -+ if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { -+ std::cerr << "ggml_vulkan: buffer_write_async dst buffer is host_visible. Use synchronous write." << std::endl; -+ GGML_ABORT("fatal error"); -+ } -+ // Check if src is pinned memory -+ vk_buffer buf = nullptr; -+ size_t buf_offset = 0; -+ ggml_vk_host_get(dst->device, src, buf, buf_offset); -+ -+ if (buf != nullptr) { -+ // Memory is pinned, use as staging buffer -+ std::vector slices(1); -+ if (width == spitch) { -+ // Only do single write if stride is equal -+ slices[0].srcOffset = buf_offset; -+ slices[0].dstOffset = offset; -+ slices[0].size = width * height; -+ } else { -+ slices.resize(height); -+ for (size_t i = 0; i < height; i++) { -+ slices[i].srcOffset = buf_offset + i * spitch; -+ slices[i].dstOffset = offset + i * width; -+ slices[i].size = width; -+ } -+ } -+ -+ ggml_vk_sync_buffers(subctx); -+ subctx->s->buffer.copyBuffer(buf->buffer, dst->buffer, slices); -+ return; -+ } -+ VK_LOG_DEBUG("STAGING"); -+ -+ if (!sync_staging) { -+ GGML_ABORT("Asynchronous write to non-pinned memory not supported"); -+ } -+ -+ // Staging buffer required -+ const size_t copy_size = width*height; -+ ggml_vk_ensure_sync_staging_buffer(dst->device, copy_size); -+ -+ vk_buffer& staging_buffer = dst->device->sync_staging; -+ -+ VkBufferCopy buf_copy = { -+ 0, -+ offset, -+ copy_size}; -+ -+ ggml_vk_sync_buffers(subctx); -+ vkCmdCopyBuffer(subctx->s->buffer, (VkBuffer)staging_buffer->buffer, (VkBuffer)dst->buffer, 1, &buf_copy); -+ -+ if (width == spitch) { -+ deferred_memcpy((uint8_t *)staging_buffer->ptr, src, width * height, &subctx->in_memcpys); -+ } else { -+ for (size_t i = 0; i < height; i++) { -+ deferred_memcpy((uint8_t *)staging_buffer->ptr + i * width, (const uint8_t *) src + i * spitch, width, &subctx->in_memcpys); -+ } -+ } -+} -+ -+static void ggml_vk_buffer_write_async(vk_context subctx, vk_buffer& dst, size_t offset, const void * src, size_t size, bool sync_staging = false) { -+ VK_LOG_DEBUG("ggml_vk_buffer_write_async(" << size << ")"); -+ return ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, size, size, 1, sync_staging); -+} -+ -+static void ggml_vk_buffer_write_2d(vk_buffer& dst, size_t offset, const void * src, size_t spitch, size_t width, size_t height) { -+ VK_LOG_DEBUG("ggml_vk_buffer_write_2d(" << width << ", " << height << ")"); -+ // Buffer is already mapped -+ if(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) { -+ GGML_ASSERT(dst->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent); -+ -+ for (size_t i = 0; i < height; i++) { -+ memcpy((uint8_t *)dst->ptr + offset + i * width, (const uint8_t *) src + i * spitch, width); -+ } -+ } else { -+ vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue); -+ ggml_vk_ctx_begin(dst->device, subctx); -+ ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, spitch, width, height, true); -+ ggml_vk_ctx_end(subctx); -+ -+ for (auto& cpy : subctx->in_memcpys) { -+ memcpy(cpy.dst, cpy.src, cpy.n); -+ } -+ -+ ggml_vk_submit(subctx, dst->device->fence); -+ VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_buffer_write_2d waitForFences"); -+ dst->device->device.resetFences({ dst->device->fence }); -+ } -+} -+ -+static void ggml_vk_buffer_write(vk_buffer& dst, size_t offset, const void * src, size_t size) { -+ VK_LOG_DEBUG("ggml_vk_buffer_write(" << size << ")"); -+ ggml_vk_buffer_write_2d(dst, offset, src, 0, size, 1); -+} -+ -+static void ggml_vk_buffer_read_2d_async(vk_context subctx, vk_buffer& src, size_t offset, void * dst, size_t spitch, size_t dpitch, size_t width, size_t height, bool sync_staging = false) { -+ VK_LOG_DEBUG("ggml_vk_buffer_read_2d_async(offset=" << offset << ", width=" << width << ", height=" << height << ")"); -+ GGML_ASSERT(width > 0); -+ GGML_ASSERT(height > 0); -+ GGML_ASSERT(src != nullptr); -+ -+ // TODO: staging_offset is not used -+ -+ // Check if dst is pinned memory -+ vk_buffer buf = nullptr; -+ size_t buf_offset = 0; -+ ggml_vk_host_get(src->device, dst, buf, buf_offset); -+ -+ std::vector slices(1); -+ if (width == spitch && width == dpitch) { -+ // Only do single write if stride is equal -+ slices[0].srcOffset = offset; -+ slices[0].dstOffset = buf_offset; -+ slices[0].size = width * height; -+ } else { -+ slices.resize(height); -+ for (size_t i = 0; i < height; i++) { -+ slices[i].srcOffset = offset + i * spitch; -+ slices[i].dstOffset = buf_offset + i * dpitch; -+ slices[i].size = width; -+ } -+ } -+ -+ if (buf != nullptr) { -+ // Memory is pinned, use as staging buffer -+ ggml_vk_sync_buffers(subctx); -+ subctx->s->buffer.copyBuffer(src->buffer, buf->buffer, slices); -+ -+ return; -+ } -+ VK_LOG_DEBUG("STAGING"); -+ -+ if (!sync_staging) { -+ GGML_ABORT("Asynchronous read from non-pinned memory not supported"); -+ } -+ -+ // Fall back to staging buffer -+ const size_t copy_size = dpitch * height; -+ ggml_vk_ensure_sync_staging_buffer(src->device, copy_size); -+ -+ vk_buffer& staging_buffer = src->device->sync_staging; -+ -+ ggml_vk_sync_buffers(subctx); -+ subctx->s->buffer.copyBuffer(src->buffer, staging_buffer->buffer, slices); -+ -+ deferred_memcpy(dst, staging_buffer->ptr, copy_size, &subctx->out_memcpys); -+} -+ -+static void ggml_vk_buffer_read_async(vk_context subctx, vk_buffer& src, size_t offset, void * dst, size_t size, bool sync_staging = false) { -+ return ggml_vk_buffer_read_2d_async(subctx, src, offset, dst, size, size, size, 1, sync_staging); -+} -+ -+static void ggml_vk_buffer_read(vk_buffer& src, size_t offset, void * dst, size_t size) { -+ VK_LOG_DEBUG("ggml_vk_buffer_read(" << src->buffer << ", " << offset << ", " << size << ")"); -+ -+ // If the device is not an UMA device the memory is host-accessible through rebar. While writing -+ // through PCIe is sufficient fast reading back data from PCIe is slower than going through -+ // the HW device to host copy path. -+ if(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible && src->device->uma) { -+ GGML_ASSERT(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent); -+ -+ memcpy(dst, (uint8_t *) src->ptr + offset, size); -+ } else { -+ vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue); -+ ggml_vk_ctx_begin(src->device, subctx); -+ ggml_vk_buffer_read_async(subctx, src, offset, dst, size, true); -+ ggml_vk_ctx_end(subctx); -+ -+ ggml_vk_submit(subctx, src->device->fence); -+ VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_read waitForFences"); -+ src->device->device.resetFences({ src->device->fence }); -+ -+ for (auto& cpy : subctx->out_memcpys) { -+ memcpy(cpy.dst, cpy.src, cpy.n); -+ } -+ } -+} -+ -+static void ggml_vk_buffer_copy_async(vk_context& ctx, vk_buffer& dst, size_t dst_offset, vk_buffer& src, size_t src_offset, size_t size) { -+ VK_LOG_DEBUG("ggml_vk_buffer_copy_async(" << size << ")"); -+ // Make sure both buffers are on same device -+ GGML_ASSERT(src->device == dst->device); -+ -+ VkBufferCopy bc{ src_offset, dst_offset, size }; -+ -+ vkCmdCopyBuffer(ctx->s->buffer, (VkBuffer)src->buffer, (VkBuffer)dst->buffer, 1, &bc); -+} -+ -+static void ggml_vk_buffer_copy(vk_buffer& dst, size_t dst_offset, vk_buffer& src, size_t src_offset, size_t size) { -+ if (src->device == dst->device) { -+ VK_LOG_DEBUG("ggml_vk_buffer_copy(SINGLE_DEVICE, " << size << ")"); -+ // Copy within the device -+ vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue); -+ ggml_vk_ctx_begin(src->device, subctx); -+ ggml_vk_buffer_copy_async(subctx, dst, dst_offset, src, src_offset, size); -+ ggml_vk_ctx_end(subctx); -+ ggml_vk_submit(subctx, src->device->fence); -+ VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_copy waitForFences"); -+ src->device->device.resetFences({ src->device->fence }); -+ } else { -+ VK_LOG_DEBUG("ggml_vk_buffer_copy(MULTI_DEVICE, " << size << ")"); -+ // Copy device to device -+ ggml_vk_ensure_sync_staging_buffer(src->device, size); -+ ggml_vk_ensure_sync_staging_buffer(dst->device, size); -+ -+ // Copy to src staging buffer -+ ggml_vk_buffer_copy(src->device->sync_staging, 0, src, src_offset, size); -+ // memcpy to dst staging buffer -+ memcpy(dst->device->sync_staging->ptr, src->device->sync_staging->ptr, size); -+ // Copy to dst buffer -+ ggml_vk_buffer_copy(dst, dst_offset, dst->device->sync_staging, 0, size); -+ } -+} -+ -+static void ggml_vk_buffer_memset(vk_buffer& dst, size_t offset, uint32_t c, size_t size) { -+ VK_LOG_DEBUG("ggml_vk_buffer_memset(" << offset << ", " << c << ", " << size << ")"); -+ -+ vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue); -+ ggml_vk_ctx_begin(dst->device, subctx); -+ subctx->s->buffer.fillBuffer(dst->buffer, offset, size, c); -+ ggml_vk_ctx_end(subctx); -+ -+ ggml_vk_submit(subctx, dst->device->fence); -+ VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_memset waitForFences"); -+ dst->device->device.resetFences({ dst->device->fence }); -+} -+ -+static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, int m, int n, int k, const vk_pipeline& pipeline) { -+ VK_LOG_DEBUG("ggml_vk_guess_split_k(" << m << ", " << n << ", " << k << ")"); -+ -+ uint32_t split_k = 1; -+ if (ctx->device->shader_core_count != 0 && m >= (int)pipeline->wg_denoms[0] && n >= (int)pipeline->wg_denoms[1]) { -+ // If k is 'large' and the SMs will fill less than halfway, use split_k. -+ uint32_t m_tiles = CEIL_DIV(m, pipeline->wg_denoms[0]); -+ uint32_t n_tiles = CEIL_DIV(n, pipeline->wg_denoms[1]); -+ if (k >= 2048 && m_tiles * n_tiles < ctx->device->shader_core_count / 2) { -+ split_k = ctx->device->shader_core_count / (m_tiles * n_tiles); -+ // Clamp to 2 or 4 -+ split_k = std::min(split_k, 4u); -+ if (split_k == 3) { -+ split_k = 2; -+ } -+ } -+ } -+ -+ return split_k; -+} -+ -+static vk_pipeline ggml_vk_guess_matmul_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, bool aligned) { -+ VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline(" << m << ", " << n << ", " << aligned << ")"); -+ -+ if (ctx->device->coopmat2) { -+ if ((ctx->device->mul_mat_l && (m % mmp->l->wg_denoms[0]) == 0 && (n % mmp->l->wg_denoms[1]) == 0) || (!ctx->device->mul_mat_m && !ctx->device->mul_mat_s)) { -+ return aligned ? mmp->a_l : mmp->l; -+ } -+ if ((ctx->device->mul_mat_m && (m % mmp->m->wg_denoms[0]) == 0 && (n % mmp->m->wg_denoms[1]) == 0) || !ctx->device->mul_mat_s) { -+ return aligned ? mmp->a_m : mmp->m; -+ } -+ return aligned ? mmp->a_s : mmp->s; -+ } -+ -+ if ((ctx->device->mul_mat_s && (m <= 32 || n <= 32)) || (!ctx->device->mul_mat_m && !ctx->device->mul_mat_l)) { -+ return aligned ? mmp->a_s : mmp->s; -+ } -+ if ((ctx->device->mul_mat_m && (m <= 64 || n <= 64)) || !ctx->device->mul_mat_l) { -+ return aligned ? mmp->a_m : mmp->m; -+ } -+ return aligned ? mmp->a_l : mmp->l; -+} -+ -+static uint32_t ggml_vk_guess_matmul_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n) { -+ VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ")"); -+ return ggml_vk_guess_matmul_pipeline(ctx, mmp, m, n, true)->align; -+} -+ -+static void ggml_vk_matmul( -+ ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, -+ vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& split_k_buffer, -+ uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, -+ uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, -+ uint32_t split_k, uint32_t batch, uint32_t ne02, uint32_t ne12, uint32_t broadcast2, uint32_t broadcast3) { -+ VK_LOG_DEBUG("ggml_vk_matmul(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), split_k: (" << (split_k_buffer.buffer != nullptr ? split_k_buffer.buffer->buffer : VK_NULL_HANDLE) << ", " << split_k_buffer.offset << ", " << split_k_buffer.size << "), m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", split_k: " << split_k << ", batch: " << batch << ", ne02: " << ne02 << ", ne12: " << ne12 << ", broadcast2: " << broadcast2 << ", broadcast3: " << broadcast3 << ")"); -+ ggml_vk_sync_buffers(subctx); -+ if (split_k == 1) { -+ const vk_mat_mat_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, k, ne02, ne12, broadcast2, broadcast3 }; -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d }, sizeof(vk_mat_mat_push_constants), &pc, { m, n, batch }); -+ return; -+ } -+ -+ GGML_ASSERT(batch_stride_d == m * n); -+ -+ const vk_mat_mat_push_constants pc1 = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, CEIL_DIV(k, split_k), ne02, ne12, broadcast2, broadcast3 }; -+ // Make sure enough workgroups get assigned for split k to work -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, split_k_buffer }, sizeof(vk_mat_mat_push_constants), &pc1, { (CEIL_DIV(m, pipeline->wg_denoms[0]) * pipeline->wg_denoms[0]) * split_k, n, batch }); -+ ggml_vk_sync_buffers(subctx); -+ const std::array pc2 = { (uint32_t)(m * n * batch), split_k }; -+ ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_matmul_split_k_reduce, { split_k_buffer, d }, pc2.size() * sizeof(uint32_t), pc2.data(), { m * n * batch, 1, 1 }); -+} -+ -+static vk_pipeline ggml_vk_guess_matmul_id_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, bool aligned) { -+ VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline(" << m << ", " << n << ", " << aligned << ")"); -+ -+ if (ctx->device->coopmat2) { -+ if ((ctx->device->mul_mat_id_l && (m % mmp->l->wg_denoms[0]) == 0 && (n % mmp->l->wg_denoms[1]) == 0) || (!ctx->device->mul_mat_id_m && !ctx->device->mul_mat_id_s)) { -+ return aligned ? mmp->a_l : mmp->l; -+ } -+ if ((ctx->device->mul_mat_id_m && (m % mmp->m->wg_denoms[0]) == 0 && (n % mmp->m->wg_denoms[1]) == 0) || !ctx->device->mul_mat_id_s) { -+ return aligned ? mmp->a_m : mmp->m; -+ } -+ return aligned ? mmp->a_s : mmp->s; -+ } -+ -+ if ((ctx->device->mul_mat_id_s && (m <= 32 || n <= 32)) || (!ctx->device->mul_mat_id_m && !ctx->device->mul_mat_id_l)) { -+ return aligned ? mmp->a_s : mmp->s; -+ } -+ if ((ctx->device->mul_mat_id_m && (m <= 64 || n <= 64)) || !ctx->device->mul_mat_id_l) { -+ return aligned ? mmp->a_m : mmp->m; -+ } -+ return aligned ? mmp->a_l : mmp->l; -+} -+ -+static uint32_t ggml_vk_guess_matmul_id_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n) { -+ VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ")"); -+ return ggml_vk_guess_matmul_id_pipeline(ctx, mmp, m, n, true)->align; -+} -+ -+static void ggml_vk_matmul_id( -+ ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, -+ vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& ids, -+ uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, -+ uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, -+ uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11) { -+ VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), " << -+ "m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " << -+ "batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " << -+ "n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")"); -+ ggml_vk_sync_buffers(subctx); -+ const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, -+ nei0, nei1, nbi1, ne11 }; -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids }, sizeof(vk_mat_mat_id_push_constants), &pc, { m, nei1, n_as }); -+} -+ -+static bool ggml_vk_dim01_contiguous(const ggml_tensor * tensor) { -+ return -+ tensor->nb[0] == ggml_type_size(tensor->type) && -+ tensor->nb[1] == (tensor->nb[0]*tensor->ne[0])/ggml_blck_size(tensor->type) && -+ tensor->nb[3] == tensor->nb[2]*tensor->ne[2]; -+} -+ -+static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src, const ggml_tensor * dst, ggml_type to) { -+ -+ // Choose "contiguous copy" shader if src/dst are contiguous -+ bool contig = ggml_is_contiguous(src) && (!dst || ggml_is_contiguous(dst)); -+ -+ if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F32) { -+ if (contig) { -+ return ctx->device->pipeline_contig_cpy_f32_f32; -+ } else { -+ return ctx->device->pipeline_cpy_f32_f32; -+ } -+ } -+ if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F16) { -+ if (contig) { -+ return ctx->device->pipeline_contig_cpy_f32_f16; -+ } else { -+ return ctx->device->pipeline_cpy_f32_f16; -+ } -+ } -+ if (src->type == GGML_TYPE_F16 && to == GGML_TYPE_F16) { -+ if (contig) { -+ return ctx->device->pipeline_contig_cpy_f16_f16; -+ } else { -+ return ctx->device->pipeline_cpy_f16_f16; -+ } -+ } -+ -+ std::cerr << "Missing CPY op for types: " << ggml_type_name(src->type) << " " << ggml_type_name(to) << std::endl; -+ GGML_ABORT("fatal error"); -+} -+ -+static void ggml_vk_cpy_to_contiguous(ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline pipeline, const ggml_tensor * tensor, vk_subbuffer&& in, vk_subbuffer&& out) { -+ VK_LOG_DEBUG("ggml_vk_cpy_to_contiguous((" << tensor << ", type=" << tensor->type << ", ne0=" << tensor->ne[0] << ", ne1=" << tensor->ne[1] << ", ne2=" << tensor->ne[2] << ", ne3=" << tensor->ne[3] << ", nb0=" << tensor->nb[0] << ", nb1=" << tensor->nb[1] << ", nb2=" << tensor->nb[2] << ", nb3=" << tensor->nb[3] << "), "; -+ std::cerr << "buffer in size=" << in.buffer->size << ", buffer out size=" << out.buffer->size << ")"); -+ const int tensor_type_size = ggml_type_size(tensor->type); -+ -+ const uint32_t ne = ggml_nelements(tensor); -+ std::array elements; -+ -+ if (ne > 262144) { -+ elements = { 512, 512, CEIL_DIV(ne, 262144) }; -+ } else if (ne > 512) { -+ elements = { 512, CEIL_DIV(ne, 512), 1 }; -+ } else { -+ elements = { ne, 1, 1 }; -+ } -+ -+ vk_op_unary_push_constants pc = { -+ (uint32_t)ne, -+ (uint32_t)tensor->ne[0], (uint32_t)tensor->ne[1], (uint32_t)tensor->ne[2], (uint32_t)tensor->ne[3], (uint32_t)tensor->nb[0] / tensor_type_size, (uint32_t)tensor->nb[1] / tensor_type_size, (uint32_t)tensor->nb[2] / tensor_type_size, (uint32_t)tensor->nb[3] / tensor_type_size, -+ (uint32_t)tensor->ne[0], (uint32_t)tensor->ne[1], (uint32_t)tensor->ne[2], (uint32_t)tensor->ne[3], 1 , (uint32_t)tensor->ne[0] , (uint32_t)(tensor->ne[0] * tensor->ne[1]) , (uint32_t)(tensor->ne[0] * tensor->ne[1] * tensor->ne[2]), -+ 0, -+ 0.0f, 0.0f, -+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -+ }; -+ init_pushconst_fastdiv(pc); -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { in, out }, sizeof(vk_op_unary_push_constants), &pc, elements); -+} -+ -+static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ VK_LOG_DEBUG("ggml_vk_mul_mat_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; -+ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; -+ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; -+ std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); -+ GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT -+ GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT -+ -+ const uint64_t ne00 = src0->ne[0]; -+ const uint64_t ne01 = src0->ne[1]; -+ const uint64_t ne02 = src0->ne[2]; -+ const uint64_t ne03 = src0->ne[3]; -+ -+ const uint64_t ne10 = src1->ne[0]; -+ const uint64_t ne11 = src1->ne[1]; -+ const uint64_t ne12 = src1->ne[2]; -+ const uint64_t ne13 = src1->ne[3]; -+ -+ const uint64_t ne20 = dst->ne[0]; -+ const uint64_t ne21 = dst->ne[1]; -+ -+ const uint64_t r2 = ne12 / ne02; -+ const uint64_t r3 = ne13 / ne03; -+ -+ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; -+ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; -+ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; -+ -+ vk_buffer d_Qx = nullptr; -+ size_t qx_buf_offset = 0; -+ vk_buffer d_Qy = nullptr; -+ size_t qy_buf_offset = 0; -+ -+ bool src0_uma = false; -+ bool src1_uma = false; -+ -+ if (ctx->device->uma) { -+ ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); -+ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); -+ src0_uma = d_Qx != nullptr; -+ src1_uma = d_Qy != nullptr; -+ } -+ -+ const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); -+ // Reformat and convert to fp16 if src1 is non-contiguous, or for coopmat2 for better perf -+ const bool y_non_contig = (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || -+ !ggml_vk_dim01_contiguous(src1); -+ -+ const bool y_f32_kernel = src1->type == GGML_TYPE_F32 && !y_non_contig; -+ -+ vk_matmul_pipeline mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, src0->type, y_non_contig ? GGML_TYPE_F16 : src1->type, (ggml_prec)dst->op_params[0]); -+ -+ const bool qx_needs_dequant = mmp == nullptr || x_non_contig; -+ const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !y_f32_kernel) || y_non_contig; -+ -+ if (qx_needs_dequant) { -+ // Fall back to dequant + f16 mulmat -+ mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, GGML_TYPE_F16, y_f32_kernel ? GGML_TYPE_F32 : GGML_TYPE_F16, (ggml_prec)dst->op_params[0]); -+ } -+ -+ // Not implemented -+ GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT -+ -+ const int x_ne = ne01 * ne00; -+ const int y_ne = ne11 * ne10; -+ const int d_ne = ne11 * ne01; -+ -+ const uint32_t kpad = ggml_vk_align_size(ne10, ggml_vk_guess_matmul_pipeline_align(ctx, mmp, ne01, ne11)); -+ const bool aligned = ne10 == kpad && ne01 > 8 && ne11 > 8; -+ -+ vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline(ctx, mmp, ne01, ne11, aligned); -+ -+ const uint32_t split_k = ggml_vk_guess_split_k(ctx, ne01, ne11, ne10, pipeline); -+ -+ const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type); -+ const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); -+ const uint64_t x_sz = !qx_needs_dequant ? qx_sz : sizeof(ggml_fp16_t) * x_ne; -+ const uint64_t y_sz = y_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; -+ const uint64_t d_sz = sizeof(float) * d_ne; -+ -+ vk_pipeline to_fp16_vk_0 = nullptr; -+ vk_pipeline to_fp16_vk_1 = nullptr; -+ -+ if (x_non_contig) { -+ to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, GGML_TYPE_F16); -+ } else { -+ to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, src0->type); -+ } -+ if (y_non_contig) { -+ to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, GGML_TYPE_F16); -+ } else { -+ to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); -+ } -+ GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT -+ GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT -+ -+ if (dryrun) { -+ const uint64_t x_sz_upd = x_sz * ne02 * ne03; -+ const uint64_t y_sz_upd = y_sz * ne12 * ne13; -+ const uint64_t split_k_size = split_k > 1 ? d_sz * ne12 * ne13 * split_k : 0; -+ if ( -+ (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || -+ (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size) || -+ (split_k > 1 && split_k_size > ctx->device->max_memory_allocation_size)) { -+ GGML_ABORT("Requested preallocation size is too large"); -+ } -+ if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { -+ ctx->prealloc_size_x = x_sz_upd; -+ } -+ if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { -+ ctx->prealloc_size_y = y_sz_upd; -+ } -+ if (split_k > 1 && ctx->prealloc_size_split_k < split_k_size) { -+ ctx->prealloc_size_split_k = split_k_size; -+ } -+ -+ // Request descriptor sets -+ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); -+ if (qx_needs_dequant) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); -+ } -+ if (qy_needs_dequant) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); -+ } -+ if (split_k > 1) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, 1); -+ } -+ return; -+ } -+ -+ vk_buffer d_D = dst_buf_ctx->dev_buffer; -+ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; -+ GGML_ASSERT(d_D != nullptr); -+ GGML_ASSERT(d_D->size >= d_buf_offset + d_sz * ne02 * ne03); -+ vk_buffer d_X; -+ uint64_t x_buf_offset = 0; -+ vk_buffer d_Y; -+ uint64_t y_buf_offset = 0; -+ if (!src0_uma) { -+ d_Qx = src0_buf_ctx->dev_buffer; -+ qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; -+ GGML_ASSERT(d_Qx != nullptr); -+ } -+ if (!src1_uma) { -+ d_Qy = src1_buf_ctx->dev_buffer; -+ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; -+ GGML_ASSERT(d_Qy != nullptr); -+ } -+ if (qx_needs_dequant) { -+ d_X = ctx->prealloc_x; -+ GGML_ASSERT(d_X->size >= x_sz * ne02 * ne03); -+ } else { -+ d_X = d_Qx; -+ x_buf_offset = qx_buf_offset; -+ GGML_ASSERT(qx_sz == x_sz); -+ } -+ if (qy_needs_dequant) { -+ d_Y = ctx->prealloc_y; -+ GGML_ASSERT(d_Y->size >= y_sz * ne02 * ne03); -+ } else { -+ d_Y = d_Qy; -+ y_buf_offset = qy_buf_offset; -+ GGML_ASSERT(qy_sz == y_sz); -+ } -+ -+ if (x_non_contig) { -+ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); -+ } else if (qx_needs_dequant) { -+ const std::vector pc = { (uint32_t)ne01, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)(ggml_nelements(src0)) }; -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz * ne02 * ne03 }, vk_subbuffer{ d_X, 0, x_sz * ne02 * ne03 } }, pc.size() * sizeof(uint32_t), pc.data(), { (uint32_t)(x_ne * ne02 * ne03), 1, 1}); -+ } -+ if (y_non_contig) { -+ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); -+ } -+ -+ uint32_t stride_batch_x = ne00*ne01; -+ uint32_t stride_batch_y = ne10*ne11; -+ -+ if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { -+ stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); -+ } -+ -+ if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { -+ stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); -+ } -+ -+ // compute -+ ggml_vk_matmul( -+ ctx, subctx, pipeline, -+ { d_X, x_buf_offset, x_sz * ne02 * ne03 }, { d_Y, y_buf_offset, y_sz * ne12 * ne13 }, -+ { d_D, d_buf_offset, d_sz * ne12 * ne13 }, { ctx->prealloc_split_k, 0, d_sz * ne12 * ne13 * split_k }, -+ ne01, ne11, ne10, -+ ne10, ne10, ne01, stride_batch_x, stride_batch_y, ne20*ne21, -+ split_k, ne12*ne13, ne02, ne12, r2, r3 -+ ); // NOLINT -+} -+ -+static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ VK_LOG_DEBUG("ggml_vk_mul_mat_vec_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; -+ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; -+ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; -+ std::cerr << "), " << (dryrun ? "dryrun" : "") << "),)"); -+ GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT -+ GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT -+ -+ const uint64_t ne00 = src0->ne[0]; -+ const uint64_t ne01 = src0->ne[1]; -+ const uint64_t ne02 = src0->ne[2]; -+ const uint64_t ne03 = src0->ne[3]; -+ -+ const uint64_t ne10 = src1->ne[0]; -+ const uint64_t ne11 = src1->ne[1]; -+ const uint64_t ne12 = src1->ne[2]; -+ const uint64_t ne13 = src1->ne[3]; -+ -+ const uint64_t ne20 = dst->ne[0]; -+ const uint64_t ne21 = dst->ne[1]; -+ const uint64_t ne22 = dst->ne[2]; -+ const uint64_t ne23 = dst->ne[3]; -+ -+ const uint64_t r2 = ne12 / ne02; -+ const uint64_t r3 = ne13 / ne03; -+ -+ // batch_n indicates that we need to compute a few vector results, and this assumes -+ // ne12 and ne13 are 1. It overloads the batch_strides to hold the row strides. -+ GGML_ASSERT(ne11 == 1 || ne12 * ne13 == 1); -+ bool batch_n = ne11 > 1; -+ -+ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; -+ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; -+ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; -+ -+ vk_buffer d_Qx = nullptr; -+ size_t qx_buf_offset = 0; -+ vk_buffer d_Qy = nullptr; -+ size_t qy_buf_offset = 0; -+ -+ bool src0_uma = false; -+ bool src1_uma = false; -+ -+ if (ctx->device->uma) { -+ ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); -+ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); -+ src0_uma = d_Qx != nullptr; -+ src1_uma = d_Qy != nullptr; -+ } -+ -+ const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); -+ const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); -+ -+ const bool f16_f32_kernel = src1->type == GGML_TYPE_F32; -+ -+ const bool qx_needs_dequant = x_non_contig; -+ const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !f16_f32_kernel) || y_non_contig; -+ -+ // Not implemented -+ GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT -+ -+ const uint64_t x_ne = ne01 * ne00; -+ const uint64_t y_ne = ne11 * ne10; -+ const uint64_t d_ne = ne11 * ne01; -+ -+ const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); -+ const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); -+ const uint64_t x_sz = x_non_contig ? ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment) : qx_sz; -+ const uint64_t y_sz = f16_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; -+ const uint64_t d_sz = sizeof(float) * d_ne; -+ -+ vk_pipeline to_fp16_vk_0 = nullptr; -+ vk_pipeline to_fp16_vk_1 = nullptr; -+ if (x_non_contig) { -+ to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, src0->type); -+ } -+ if (y_non_contig) { -+ to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, src1->type); -+ } else { -+ to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); -+ } -+ vk_pipeline dmmv = ggml_vk_get_dequantize_mul_mat_vec(ctx, src0->type, src1->type, ne11); -+ GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT -+ GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT -+ GGML_ASSERT(dmmv != nullptr); -+ -+ if (dryrun) { -+ const uint64_t x_sz_upd = x_sz * ne02 * ne03; -+ const uint64_t y_sz_upd = y_sz * ne12 * ne13; -+ if ( -+ (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || -+ (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { -+ GGML_ABORT("Requested preallocation size is too large"); -+ } -+ if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { -+ ctx->prealloc_size_x = x_sz_upd; -+ } -+ if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { -+ ctx->prealloc_size_y = y_sz_upd; -+ } -+ -+ // Request descriptor sets -+ if (qx_needs_dequant) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); -+ } -+ if (qy_needs_dequant) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); -+ } -+ ggml_pipeline_request_descriptor_sets(ctx->device, dmmv, 1); -+ return; -+ } -+ -+ vk_buffer d_D = dst_buf_ctx->dev_buffer; -+ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; -+ GGML_ASSERT(d_D != nullptr); -+ vk_buffer d_X; -+ uint64_t x_buf_offset = 0; -+ vk_buffer d_Y; -+ uint64_t y_buf_offset = 0; -+ if(!src0_uma) { -+ d_Qx = src0_buf_ctx->dev_buffer; -+ qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; -+ GGML_ASSERT(d_Qx != nullptr); -+ } -+ if(!src1_uma) { -+ d_Qy = src1_buf_ctx->dev_buffer; -+ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; -+ GGML_ASSERT(d_Qy != nullptr); -+ } -+ if (qx_needs_dequant) { -+ d_X = ctx->prealloc_x; -+ } else { -+ d_X = d_Qx; -+ x_buf_offset = qx_buf_offset; -+ GGML_ASSERT(qx_sz == x_sz); -+ } -+ if (qy_needs_dequant) { -+ d_Y = ctx->prealloc_y; -+ } else { -+ d_Y = d_Qy; -+ y_buf_offset = qy_buf_offset; -+ GGML_ASSERT(qy_sz == y_sz); -+ } -+ -+ if (x_non_contig) { -+ GGML_ASSERT(x_sz == ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment)); -+ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); -+ } -+ if (y_non_contig) { -+ GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); -+ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); -+ } -+ -+ // For batch_n, the A matrix is the same for each batch, and B/D use the row stride as the batch stride -+ uint32_t stride_batch_x = batch_n ? 0 : ne00*ne01; -+ uint32_t stride_batch_y = batch_n ? ne10 : (ne10*ne11); -+ uint32_t stride_batch_d = batch_n ? ne20 : (ne20*ne21); -+ -+ if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { -+ stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); -+ } -+ -+ if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { -+ stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); -+ } -+ -+ const uint32_t max_groups_x = ctx->device->properties.limits.maxComputeWorkGroupCount[0]; -+ -+ uint32_t groups_x = ne01; -+ uint32_t groups_z = 1; -+ -+ if (ne01 > max_groups_x) { -+ groups_z = 64; -+ groups_x = CEIL_DIV(groups_x, groups_z); -+ } -+ -+ // compute -+ const vk_mat_vec_push_constants pc = { -+ (uint32_t)ne00, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne01, -+ stride_batch_x, stride_batch_y, stride_batch_d, -+ (uint32_t)ne02, (uint32_t)ne12, (uint32_t)r2, (uint32_t)r3, -+ }; -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, dmmv, -+ { vk_subbuffer{ d_X, x_buf_offset, x_sz * ne02 * ne03 }, vk_subbuffer{ d_Y, y_buf_offset, y_sz * ne12 * ne13 }, vk_subbuffer{ d_D, d_buf_offset, d_sz * ne22 * ne23} }, -+ sizeof(vk_mat_vec_push_constants), &pc, { groups_x, (uint32_t)(ne12 * ne13), groups_z }); -+} -+ -+static void ggml_vk_mul_mat_vec_p021_f16_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ VK_LOG_DEBUG("ggml_vk_mul_mat_p021_f16_f32(" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; -+ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; -+ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; -+ std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); -+ GGML_ASSERT(ggml_is_permuted(src0) && ggml_is_permuted(src1)); -+ GGML_ASSERT(src0->nb[0] <= src0->nb[1] && src0->nb[2] <= src0->nb[3]); // NOLINT -+ GGML_ASSERT(src1->nb[0] <= src1->nb[1] && src1->nb[2] <= src1->nb[3]); // NOLINT -+ GGML_ASSERT(src0->type == GGML_TYPE_F16); -+ GGML_ASSERT(src1->type == GGML_TYPE_F32); -+ -+ const uint64_t ne00 = src0->ne[0]; -+ const uint64_t ne01 = src0->ne[1]; -+ const uint64_t ne02 = src0->ne[2]; -+ // const uint64_t ne03 = src0->ne[3]; -+ -+ const uint64_t ne10 = src1->ne[0]; -+ const uint64_t ne11 = src1->ne[1]; -+ const uint64_t ne12 = src1->ne[2]; -+ // const uint64_t ne13 = src1->ne[3]; -+ -+ GGML_ASSERT(ne11 == 1); -+ -+ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; -+ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; -+ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; -+ -+ vk_buffer d_Qy = nullptr; -+ size_t qy_buf_offset = 0; -+ -+ bool src1_uma = false; -+ -+ if (ctx->device->uma) { -+ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); -+ src1_uma = d_Qy != nullptr; -+ } -+ -+ const uint64_t x_ne = ne00 * ne01 * ne02; -+ const uint64_t y_ne = ne10 * ne11 * ne12; -+ const uint64_t d_ne = ne01 * ne11 * ne12; -+ -+ const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); -+ const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); -+ const uint64_t d_sz = sizeof(float) * d_ne; -+ -+ if (dryrun) { -+ // Request descriptor sets -+ ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_mul_mat_vec_p021_f16_f32, 1); -+ return; -+ } -+ -+ vk_buffer d_D = dst_buf_ctx->dev_buffer; -+ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; -+ GGML_ASSERT(d_D != nullptr); -+ vk_buffer d_Qx = src0_buf_ctx->dev_buffer; -+ const uint64_t qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; -+ GGML_ASSERT(d_Qx != nullptr); -+ if (!src1_uma) { -+ d_Qy = src1_buf_ctx->dev_buffer; -+ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; -+ GGML_ASSERT(d_Qx != nullptr); -+ } -+ -+ const uint64_t qy_buffer_offset = (qy_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; -+ const uint64_t qy_shader_offset = qy_buf_offset - qy_buffer_offset; -+ -+ const uint64_t d_buffer_offset = (d_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; -+ const uint64_t d_shader_offset = d_buf_offset - d_buffer_offset; -+ -+ // compute -+ const std::array pc = { (uint32_t)ne00, (uint32_t)ne01, (uint32_t)ne02, (uint32_t)ne12, (uint32_t)(qy_shader_offset / ggml_type_size(src1->type)), (uint32_t)(d_shader_offset / ggml_type_size(dst->type)) }; -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_mul_mat_vec_p021_f16_f32, { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz }, vk_subbuffer{ d_Qy, qy_buffer_offset, qy_sz + qy_shader_offset }, vk_subbuffer{ d_D, d_buffer_offset, d_sz + d_shader_offset } }, 6 * sizeof(uint32_t), &pc, { 1, (uint32_t)ne01, (uint32_t)ne12 }); -+} -+ -+static void ggml_vk_mul_mat_vec_nc_f16_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ VK_LOG_DEBUG("ggml_vk_mul_mat_nc_f16_f32((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; -+ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; -+ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; -+ std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); -+ GGML_ASSERT(!ggml_is_transposed(src0)); -+ GGML_ASSERT(!ggml_is_transposed(src1)); -+ GGML_ASSERT(!ggml_is_permuted(src0)); -+ GGML_ASSERT(src0->type == GGML_TYPE_F16); -+ GGML_ASSERT(src1->type == GGML_TYPE_F32); -+ -+ const uint64_t ne00 = src0->ne[0]; -+ const uint64_t ne01 = src0->ne[1]; -+ const uint64_t ne02 = src0->ne[2]; -+ // const uint64_t ne03 = src0->ne[3]; -+ -+ const uint64_t nb01 = src0->nb[1]; -+ const uint64_t nb02 = src0->nb[2]; -+ -+ // const uint64_t ne10 = src1->ne[0]; -+ const uint64_t ne11 = src1->ne[1]; -+ const uint64_t ne12 = src1->ne[2]; -+ // const uint64_t ne13 = src1->ne[3]; -+ -+ GGML_ASSERT(ne11 == 1); -+ -+ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; -+ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; -+ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; -+ -+ vk_buffer d_Qy = nullptr; -+ size_t qy_buf_offset = 0; -+ -+ bool src1_uma = false; -+ -+ if (ctx->device->uma) { -+ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); -+ src1_uma = d_Qy != nullptr; -+ } -+ -+ const uint64_t d_ne = ne01 * ne11 * ne12; -+ -+ const uint32_t row_stride_x = nb01 / sizeof(ggml_fp16_t); -+ const uint32_t channel_stride_x = nb02 / sizeof(ggml_fp16_t); -+ -+ const uint64_t qx_sz = ggml_nbytes(src0); -+ const uint64_t qy_sz = ggml_nbytes(src1); -+ const uint64_t d_sz = sizeof(float) * d_ne; -+ -+ if (dryrun) { -+ // Request descriptor sets -+ ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_mul_mat_vec_nc_f16_f32, 1); -+ return; -+ } -+ -+ vk_buffer d_D = dst_buf_ctx->dev_buffer; -+ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; -+ GGML_ASSERT(d_D != nullptr); -+ vk_buffer d_Qx = src0_buf_ctx->dev_buffer; -+ const uint64_t qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; -+ GGML_ASSERT(d_Qx != nullptr); -+ if (!src1_uma) { -+ d_Qy = src1_buf_ctx->dev_buffer; -+ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; -+ GGML_ASSERT(d_Qx != nullptr); -+ } -+ -+ const uint64_t qy_buffer_offset = (qy_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; -+ const uint64_t qy_shader_offset = qy_buf_offset - qy_buffer_offset; -+ -+ const uint64_t d_buffer_offset = (d_buf_offset / ctx->device->properties.limits.minStorageBufferOffsetAlignment) * ctx->device->properties.limits.minStorageBufferOffsetAlignment; -+ const uint64_t d_shader_offset = d_buf_offset - d_buffer_offset; -+ -+ // compute -+ const std::array pc = { (uint32_t)ne00, (uint32_t)ne01, row_stride_x, channel_stride_x, (uint32_t)(ne12 / ne02), (uint32_t)(qy_shader_offset / ggml_type_size(src1->type)), (uint32_t)(d_shader_offset / ggml_type_size(dst->type)) }; -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_mul_mat_vec_nc_f16_f32, -+ { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz }, vk_subbuffer{ d_Qy, qy_buffer_offset, qy_sz + qy_shader_offset }, vk_subbuffer{ d_D, d_buffer_offset, d_sz + d_shader_offset } }, 7 * sizeof(uint32_t), &pc, { 1, (uint32_t)ne01, (uint32_t)ne12 }); -+} -+ -+static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ VK_LOG_DEBUG("ggml_vk_mul_mat(" << src0 << ", " << src1 << ", " << dst << ")"); -+ if (src0->type == GGML_TYPE_F16 && ggml_is_permuted(src0) && ggml_is_permuted(src1) && dst->ne[1] == 1 && -+ // detect 0213 permutation, and batch size of 1 -+ src0->nb[0] <= src0->nb[2] && -+ src0->nb[2] <= src0->nb[1] && -+ src0->nb[1] <= src0->nb[3] && -+ src1->nb[0] <= src1->nb[2] && -+ src1->nb[2] <= src1->nb[1] && -+ src1->nb[1] <= src1->nb[3] && -+ src0->ne[3] == 1 && -+ src1->ne[3] == 1) { -+ ggml_vk_mul_mat_vec_p021_f16_f32(ctx, subctx, src0, src1, dst, dryrun); -+ } else if (src0->type == GGML_TYPE_F16 && !ggml_is_contiguous(src0) && !ggml_is_transposed(src1) && dst->ne[1] == 1 && -+ !ggml_is_permuted(src0) && !ggml_is_permuted(src1)) { -+ ggml_vk_mul_mat_vec_nc_f16_f32(ctx, subctx, src0, src1, dst, dryrun); -+ // mul_mat_vec supports batching ne12*ne13 when ne11==1, or treating ne11 as the batch size (up to four) -+ // when ne12 and ne13 are one. -+ } else if ((dst->ne[1] == 1 || (dst->ne[1] <= mul_mat_vec_max_cols && src1->ne[2] * src1->ne[3] == 1)) && -+ (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type))) { -+ ggml_vk_mul_mat_vec_q_f16(ctx, subctx, src0, src1, dst, dryrun); -+ } else { -+ ggml_vk_mul_mat_q_f16(ctx, subctx, src0, src1, dst, dryrun); -+ } -+} -+ -+static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, bool dryrun = false) { -+ VK_LOG_DEBUG("ggml_vk_mul_mat_id_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; -+ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; -+ std::cerr << "), (" << ids << ", name=" << ids->name << ", type=" << ids->type << ", ne0=" << ids->ne[0] << ", ne1=" << ids->ne[1] << ", ne2=" << ids->ne[2] << ", ne3=" << ids->ne[3] << ", nb0=" << ids->nb[0] << ", nb1=" << ids->nb[1] << ", nb2=" << ids->nb[2] << ", nb3=" << ids->nb[3]; -+ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3] << "),)"); -+ GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT -+ GGML_ASSERT(ids->type == GGML_TYPE_I32); -+ -+ const uint64_t ne00 = src0->ne[0]; -+ const uint64_t ne01 = src0->ne[1]; -+ const uint64_t ne02 = src0->ne[2]; -+ const uint64_t ne03 = src0->ne[3]; -+ -+ const uint64_t ne10 = src1->ne[0]; -+ const uint64_t ne11 = src1->ne[1]; -+ const uint64_t ne12 = src1->ne[2]; -+ const uint64_t ne13 = src1->ne[3]; -+ -+ const uint64_t nei0 = ids->ne[0]; -+ const uint64_t nei1 = ids->ne[1]; -+ GGML_ASSERT(nei0 * nei1 <= 3072); -+ -+ const uint32_t nbi1 = ids->nb[1]; -+ const uint32_t nbi2 = ids->nb[2]; -+ -+ const uint64_t ne20 = dst->ne[0]; -+ const uint64_t ne21 = dst->ne[1]; -+ const uint64_t ne22 = dst->ne[2]; -+ const uint64_t ne23 = dst->ne[3]; -+ -+ const uint64_t n_as = ne02; -+ -+ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; -+ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; -+ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; -+ ggml_backend_vk_buffer_context * ids_buf_ctx = (ggml_backend_vk_buffer_context *)ids->buffer->context; -+ -+ vk_buffer d_Qx = nullptr; -+ size_t qx_buf_offset = 0; -+ vk_buffer d_Qy = nullptr; -+ size_t qy_buf_offset = 0; -+ vk_buffer d_ids = nullptr; -+ size_t ids_buf_offset = 0; -+ -+ bool src0_uma = false; -+ bool src1_uma = false; -+ bool ids_uma = false; -+ -+ if (ctx->device->uma) { -+ ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); -+ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); -+ ggml_vk_host_get(ctx->device, ids->data, d_ids, ids_buf_offset); -+ src0_uma = d_Qx != nullptr; -+ src1_uma = d_Qy != nullptr; -+ ids_uma = d_ids != nullptr; -+ } -+ -+ const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); -+ const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); -+ -+ const bool y_f32_kernel = src1->type == GGML_TYPE_F32 && !y_non_contig; -+ -+ vk_matmul_pipeline mmp = ggml_vk_get_mul_mat_mat_id_pipeline(ctx, src0->type, y_non_contig ? GGML_TYPE_F16 : src1->type, (ggml_prec)dst->op_params[0]); -+ -+ const bool qx_needs_dequant = mmp == nullptr || x_non_contig; -+ const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !y_f32_kernel) || y_non_contig; -+ -+ if (qx_needs_dequant) { -+ GGML_ABORT("fatal error"); -+ } -+ -+ // Not implemented -+ GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT -+ -+ const uint64_t x_ne = ne01 * ne00; -+ const uint64_t y_ne = ne11 * ne10; -+ const uint64_t d_ne = ne21 * ne20; -+ -+ const uint32_t kpad = ggml_vk_align_size(ne10, ggml_vk_guess_matmul_id_pipeline_align(ctx, mmp, ne01, nei1)); -+ const bool aligned = ne10 == kpad && ne01 > 8 && nei1 > 8; -+ -+ vk_pipeline pipeline = ggml_vk_guess_matmul_id_pipeline(ctx, mmp, ne01, nei1, aligned); -+ -+ const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type); -+ const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); -+ const uint64_t x_sz = !qx_needs_dequant ? qx_sz : sizeof(ggml_fp16_t) * x_ne; -+ const uint64_t y_sz = y_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; -+ const uint64_t ids_sz = nbi2; -+ const uint64_t d_sz = sizeof(float) * d_ne; -+ -+ vk_pipeline to_fp16_vk_0 = nullptr; -+ vk_pipeline to_fp16_vk_1 = nullptr; -+ -+ if (x_non_contig) { -+ to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, GGML_TYPE_F16); -+ } else { -+ to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, src0->type); -+ } -+ if (y_non_contig) { -+ to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, GGML_TYPE_F16); -+ } else { -+ to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); -+ } -+ GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT -+ GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT -+ -+ if (dryrun) { -+ const uint64_t x_sz_upd = x_sz * ne02 * ne03; -+ const uint64_t y_sz_upd = y_sz * ne12 * ne13; -+ if ( -+ (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || -+ (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { -+ GGML_ABORT("Requested preallocation size is too large"); -+ } -+ if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { -+ ctx->prealloc_size_x = x_sz_upd; -+ } -+ if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { -+ ctx->prealloc_size_y = y_sz_upd; -+ } -+ -+ // Request descriptor sets -+ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); -+ if (qx_needs_dequant) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); -+ } -+ if (qy_needs_dequant) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); -+ } -+ return; -+ } -+ -+ vk_buffer d_D = dst_buf_ctx->dev_buffer; -+ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; -+ GGML_ASSERT(d_D != nullptr); -+ vk_buffer d_X; -+ uint64_t x_buf_offset = 0; -+ vk_buffer d_Y; -+ uint64_t y_buf_offset = 0; -+ if (!src0_uma) { -+ d_Qx = src0_buf_ctx->dev_buffer; -+ qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; -+ GGML_ASSERT(d_Qx != nullptr); -+ } -+ if (!src1_uma) { -+ d_Qy = src1_buf_ctx->dev_buffer; -+ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; -+ GGML_ASSERT(d_Qy != nullptr); -+ } -+ if (!ids_uma) { -+ d_ids = ids_buf_ctx->dev_buffer; -+ ids_buf_offset = vk_tensor_offset(ids) + ids->view_offs; -+ GGML_ASSERT(d_ids != nullptr); -+ } -+ if (qx_needs_dequant) { -+ d_X = ctx->prealloc_x; -+ GGML_ASSERT(d_X->size >= x_sz * ne02 * ne03); -+ } else { -+ d_X = d_Qx; -+ x_buf_offset = qx_buf_offset; -+ GGML_ASSERT(qx_sz == x_sz); -+ } -+ if (qy_needs_dequant) { -+ d_Y = ctx->prealloc_y; -+ GGML_ASSERT(d_Y->size >= y_sz * ne02 * ne03); -+ } else { -+ d_Y = d_Qy; -+ y_buf_offset = qy_buf_offset; -+ GGML_ASSERT(qy_sz == y_sz); -+ } -+ -+ if (x_non_contig) { -+ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); -+ } else if (qx_needs_dequant) { -+ const std::vector pc = { (uint32_t)ne01, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)(ggml_nelements(src0)) }; -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, -+ { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz * ne02 * ne03 }, vk_subbuffer{ d_X, 0, x_sz * ne02 * ne03 } }, pc.size() * sizeof(uint32_t), pc.data(), { (uint32_t)(x_ne * ne02 * ne03), 1, 1}); -+ } -+ if (y_non_contig) { -+ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); -+ } -+ -+ uint32_t stride_batch_x = ne00*ne01; -+ uint32_t stride_batch_y = ne10*ne11; -+ -+ if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { -+ stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); -+ } -+ -+ if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { -+ stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); -+ } -+ -+ // compute -+ ggml_vk_matmul_id( -+ ctx, subctx, pipeline, -+ { d_X, x_buf_offset, x_sz * ne02 * ne03 }, { d_Y, y_buf_offset, y_sz * ne12 * ne13 }, -+ { d_D, d_buf_offset, d_sz * ne22 * ne23 }, { d_ids, ids_buf_offset, ids_sz }, -+ ne01, ne21, ne10, ne10, ne10, ne01, -+ stride_batch_x, stride_batch_y, ne20*ne21, -+ n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11 -+ ); // NOLINT -+} -+ -+static void ggml_vk_mul_mat_vec_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, bool dryrun = false) { -+ VK_LOG_DEBUG("ggml_vk_mul_mat_vec_id_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; -+ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; -+ std::cerr << "), (" << ids << ", name=" << ids->name << ", type=" << ids->type << ", ne0=" << ids->ne[0] << ", ne1=" << ids->ne[1] << ", ne2=" << ids->ne[2] << ", ne3=" << ids->ne[3] << ", nb0=" << ids->nb[0] << ", nb1=" << ids->nb[1] << ", nb2=" << ids->nb[2] << ", nb3=" << ids->nb[3]; -+ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; -+ std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); -+ GGML_ASSERT(ggml_vk_dim01_contiguous(src0) || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); // NOLINT -+ GGML_ASSERT(ggml_vk_dim01_contiguous(src1) || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); // NOLINT -+ GGML_ASSERT(ids->type == GGML_TYPE_I32); -+ -+ const uint64_t ne00 = src0->ne[0]; -+ const uint64_t ne01 = src0->ne[1]; -+ const uint64_t ne02 = src0->ne[2]; -+ const uint64_t ne03 = src0->ne[3]; -+ -+ const uint64_t ne10 = src1->ne[0]; -+ const uint64_t ne11 = src1->ne[1]; -+ const uint64_t ne12 = src1->ne[2]; -+ const uint64_t ne13 = src1->ne[3]; -+ -+ const uint64_t nei0 = ids->ne[0]; -+ const uint64_t nei1 = ids->ne[1]; -+ -+ const uint64_t nbi2 = ids->nb[2]; -+ -+ GGML_ASSERT(nei1 == 1); -+ -+ const uint64_t ne20 = dst->ne[0]; -+ const uint64_t ne21 = dst->ne[1]; -+ const uint64_t ne22 = dst->ne[2]; -+ const uint64_t ne23 = dst->ne[3]; -+ -+ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; -+ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; -+ ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; -+ ggml_backend_vk_buffer_context * ids_buf_ctx = (ggml_backend_vk_buffer_context *)ids->buffer->context; -+ -+ vk_buffer d_Qx = nullptr; -+ size_t qx_buf_offset = 0; -+ vk_buffer d_Qy = nullptr; -+ size_t qy_buf_offset = 0; -+ vk_buffer d_ids = nullptr; -+ size_t ids_buf_offset = 0; -+ -+ bool src0_uma = false; -+ bool src1_uma = false; -+ bool ids_uma = false; -+ -+ if (ctx->device->uma) { -+ ggml_vk_host_get(ctx->device, src0->data, d_Qx, qx_buf_offset); -+ ggml_vk_host_get(ctx->device, src1->data, d_Qy, qy_buf_offset); -+ ggml_vk_host_get(ctx->device, ids->data, d_ids, ids_buf_offset); -+ src0_uma = d_Qx != nullptr; -+ src1_uma = d_Qy != nullptr; -+ ids_uma = d_ids != nullptr; -+ } -+ -+ const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); -+ const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); -+ -+ const bool f16_f32_kernel = src1->type == GGML_TYPE_F32; -+ -+ const bool qx_needs_dequant = x_non_contig; -+ const bool qy_needs_dequant = (src1->type != GGML_TYPE_F16 && !f16_f32_kernel) || y_non_contig; -+ -+ // Not implemented -+ GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT -+ -+ const uint64_t x_ne = ne01 * ne00; -+ const uint64_t y_ne = ne11 * ne10; -+ const uint64_t d_ne = ne21 * ne20; -+ -+ const uint64_t qx_sz = ggml_vk_align_size(ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type), ctx->device->properties.limits.minStorageBufferOffsetAlignment); -+ const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type); -+ const uint64_t x_sz = x_non_contig ? ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment) : qx_sz; -+ const uint64_t y_sz = f16_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne; -+ const uint64_t ids_sz = nbi2; -+ const uint64_t d_sz = sizeof(float) * d_ne; -+ -+ vk_pipeline to_fp16_vk_0 = nullptr; -+ vk_pipeline to_fp16_vk_1 = nullptr; -+ if (x_non_contig) { -+ to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, src0, nullptr, src0->type); -+ } -+ if (y_non_contig) { -+ to_fp16_vk_1 = ggml_vk_get_cpy_pipeline(ctx, src1, nullptr, src1->type); -+ } else { -+ to_fp16_vk_1 = ggml_vk_get_to_fp16(ctx, src1->type); -+ } -+ vk_pipeline dmmv = ggml_vk_get_dequantize_mul_mat_vec_id(ctx, src0->type, src1->type); -+ GGML_ASSERT(!qx_needs_dequant || to_fp16_vk_0 != nullptr); // NOLINT -+ GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT -+ GGML_ASSERT(dmmv != nullptr); -+ -+ if (dryrun) { -+ const uint64_t x_sz_upd = x_sz * ne02 * ne03; -+ const uint64_t y_sz_upd = y_sz * ne12 * ne13; -+ if ( -+ (qx_needs_dequant && x_sz_upd > ctx->device->max_memory_allocation_size) || -+ (qy_needs_dequant && y_sz_upd > ctx->device->max_memory_allocation_size)) { -+ GGML_ABORT("Requested preallocation size is too large"); -+ } -+ if (qx_needs_dequant && ctx->prealloc_size_x < x_sz_upd) { -+ ctx->prealloc_size_x = x_sz_upd; -+ } -+ if (qy_needs_dequant && ctx->prealloc_size_y < y_sz_upd) { -+ ctx->prealloc_size_y = y_sz_upd; -+ } -+ -+ // Request descriptor sets -+ if (qx_needs_dequant) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_0, 1); -+ } -+ if (qy_needs_dequant) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, to_fp16_vk_1, 1); -+ } -+ ggml_pipeline_request_descriptor_sets(ctx->device, dmmv, 1); -+ return; -+ } -+ -+ vk_buffer d_D = dst_buf_ctx->dev_buffer; -+ const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; -+ GGML_ASSERT(d_D != nullptr); -+ vk_buffer d_X; -+ uint64_t x_buf_offset = 0; -+ vk_buffer d_Y; -+ uint64_t y_buf_offset = 0; -+ if(!src0_uma) { -+ d_Qx = src0_buf_ctx->dev_buffer; -+ qx_buf_offset = vk_tensor_offset(src0) + src0->view_offs; -+ GGML_ASSERT(d_Qx != nullptr); -+ } -+ if(!src1_uma) { -+ d_Qy = src1_buf_ctx->dev_buffer; -+ qy_buf_offset = vk_tensor_offset(src1) + src1->view_offs; -+ GGML_ASSERT(d_Qy != nullptr); -+ } -+ if(!ids_uma) { -+ d_ids = ids_buf_ctx->dev_buffer; -+ ids_buf_offset = vk_tensor_offset(ids) + ids->view_offs; -+ GGML_ASSERT(d_ids != nullptr); -+ } -+ if (qx_needs_dequant) { -+ d_X = ctx->prealloc_x; -+ } else { -+ d_X = d_Qx; -+ x_buf_offset = qx_buf_offset; -+ GGML_ASSERT(qx_sz == x_sz); -+ } -+ if (qy_needs_dequant) { -+ d_Y = ctx->prealloc_y; -+ } else { -+ d_Y = d_Qy; -+ y_buf_offset = qy_buf_offset; -+ GGML_ASSERT(qy_sz == y_sz); -+ } -+ -+ if (x_non_contig) { -+ GGML_ASSERT(x_sz == ggml_vk_align_size(ggml_type_size(src0->type) * x_ne, ctx->device->properties.limits.minStorageBufferOffsetAlignment)); -+ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, src0, { d_Qx, qx_buf_offset, VK_WHOLE_SIZE }, { d_X, 0, VK_WHOLE_SIZE }); -+ } -+ if (y_non_contig) { -+ GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); -+ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, { d_Qy, qy_buf_offset, VK_WHOLE_SIZE }, { d_Y, 0, VK_WHOLE_SIZE }); -+ } -+ -+ uint32_t stride_batch_y = ne10*ne11; -+ -+ if (!ggml_vk_dim01_contiguous(src1) && !qy_needs_dequant) { -+ stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); -+ } -+ -+ const uint32_t max_groups_x = ctx->device->properties.limits.maxComputeWorkGroupCount[0]; -+ -+ uint32_t groups_x = ne01; -+ uint32_t groups_z = 1; -+ -+ if (ne01 > max_groups_x) { -+ groups_z = 64; -+ groups_x = CEIL_DIV(groups_x, groups_z); -+ } -+ -+ // compute -+ const vk_mat_vec_id_push_constants pc = { -+ (uint32_t)ne00, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne01, -+ (uint32_t)x_ne, stride_batch_y, (uint32_t)(ne20*ne21), -+ (uint32_t)nei0, (uint32_t)ne11, -+ }; -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, dmmv, -+ { vk_subbuffer{ d_X, x_buf_offset, x_sz * ne02 * ne03 }, -+ vk_subbuffer{ d_Y, y_buf_offset, y_sz * ne12 * ne13 }, vk_subbuffer{ d_D, d_buf_offset, d_sz * ne22 * ne23}, vk_subbuffer{ d_ids, ids_buf_offset, ids_sz } }, -+ sizeof(vk_mat_vec_id_push_constants), &pc, { groups_x, (uint32_t)nei0, groups_z }); -+} -+ -+static void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, bool dryrun = false) { -+ VK_LOG_DEBUG("ggml_vk_mul_mat_id(" << src0 << ", " << src1 << ", " << src2 << ", " << dst << ")"); -+ if (src2->ne[1] == 1 && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type))) { -+ ggml_vk_mul_mat_vec_id_q_f16(ctx, subctx, src0, src1, src2, dst, dryrun); -+ } else { -+ ggml_vk_mul_mat_id_q_f16(ctx, subctx, src0, src1, src2, dst, dryrun); -+ } -+} -+ -+static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, ggml_tensor * dst, bool dryrun = false) { -+ VK_LOG_DEBUG("ggml_vk_flash_attn((" << q << ", name=" << q->name << ", type=" << q->type << ", ne0=" << q->ne[0] << ", ne1=" << q->ne[1] << ", ne2=" << q->ne[2] << ", ne3=" << q->ne[3] << ", nb0=" << q->nb[0] << ", nb1=" << q->nb[1] << ", nb2=" << q->nb[2] << ", nb3=" << q->nb[3]; -+ std::cerr << "), (" << k << ", name=" << k->name << ", type=" << k->type << ", ne0=" << k->ne[0] << ", ne1=" << k->ne[1] << ", ne2=" << k->ne[2] << ", ne3=" << k->ne[3] << ", nb0=" << k->nb[0] << ", nb1=" << k->nb[1] << ", nb2=" << k->nb[2] << ", nb3=" << k->nb[3]; -+ std::cerr << "), (" << v << ", name=" << v->name << ", type=" << v->type << ", ne0=" << v->ne[0] << ", ne1=" << v->ne[1] << ", ne2=" << v->ne[2] << ", ne3=" << v->ne[3] << ", nb0=" << v->nb[0] << ", nb1=" << v->nb[1] << ", nb2=" << v->nb[2] << ", nb3=" << v->nb[3]; -+ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; -+ std::cerr << "), " << (dryrun ? "dryrun" : "") << ")"); -+ -+ GGML_TENSOR_LOCALS(int64_t, neq, q, ne) -+ GGML_TENSOR_LOCALS(size_t, nbq, q, nb) -+ GGML_TENSOR_LOCALS(int64_t, nek, k, ne) -+ GGML_TENSOR_LOCALS(size_t, nbk, k, nb) -+ GGML_TENSOR_LOCALS(int64_t, nev, v, ne) -+ GGML_TENSOR_LOCALS(size_t, nbv, v, nb) -+ GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) -+ GGML_TENSOR_LOCALS(size_t, nb, dst, nb) -+ -+ const uint32_t nem1 = mask ? mask->ne[1] : 0; -+ const uint32_t nbm1 = mask ? mask->nb[1] : 0; -+ -+ const uint32_t D = neq0; -+ const uint32_t N = neq1; -+ const uint32_t KV = nek1; -+ -+ GGML_ASSERT(ne0 == D); -+ GGML_ASSERT(ne2 == N); -+ -+ // input tensor rows must be contiguous -+ GGML_ASSERT(nbq0 == ggml_type_size(q->type)); -+ GGML_ASSERT(nbk0 == ggml_type_size(k->type)); -+ GGML_ASSERT(nbv0 == ggml_type_size(v->type)); -+ -+ GGML_ASSERT(neq0 == D); -+ GGML_ASSERT(nek0 == D); -+ GGML_ASSERT(nev0 == D); -+ -+ GGML_ASSERT(neq1 == N); -+ GGML_ASSERT(nev0 == D); -+ -+ GGML_ASSERT(nev1 == nek1); -+ -+ // dst cannot be transposed or permuted -+ GGML_ASSERT(nb0 == sizeof(float)); -+ GGML_ASSERT(nb0 <= nb1); -+ GGML_ASSERT(nb1 <= nb2); -+ GGML_ASSERT(nb2 <= nb3); -+ -+ assert(dst->type == GGML_TYPE_F32); -+ assert(q->type == GGML_TYPE_F32); -+ assert(k->type == v->type); -+ -+ vk_pipeline *pipelines; -+ // XXX TODO other backends may be changing accumulator precision to default to f32 soon -+ bool f32acc = dst->op_params[3] == GGML_PREC_F32; -+ bool small_rows = N <= flash_attention_num_small_rows; -+ switch (D) { -+ case 64: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D64[k->type][f32acc][small_rows][0]; break; -+ case 80: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D80[k->type][f32acc][small_rows][0]; break; -+ case 96: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D96[k->type][f32acc][small_rows][0]; break; -+ case 112: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D112[k->type][f32acc][small_rows][0]; break; -+ case 128: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D128[k->type][f32acc][small_rows][0]; break; -+ case 256: pipelines = &ctx->device->pipeline_flash_attn_f32_f16_D256[k->type][f32acc][small_rows][0]; break; -+ default: -+ assert(!"unsupported D value"); -+ return; -+ } -+ assert(pipelines); -+ -+ bool aligned = (KV % pipelines[1]->align) == 0; -+ vk_pipeline pipeline = pipelines[aligned]; -+ assert(pipeline); -+ -+ if (dryrun) { -+ // Request descriptor sets -+ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); -+ return; -+ } -+ -+ float scale = 1.0f; -+ float max_bias = 0.0f; -+ float logit_softcap = 0.0f; -+ -+ memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); -+ memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); -+ memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); -+ -+ if (logit_softcap != 0) { -+ scale /= logit_softcap; -+ } -+ -+ const uint32_t n_head_kv = neq2; -+ const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head_kv)); -+ const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); -+ const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); -+ -+ ggml_vk_sync_buffers(subctx); -+ -+ vk_buffer d_Q = nullptr, d_K = nullptr, d_V = nullptr, d_D = nullptr, d_M = nullptr; -+ size_t q_buf_offset = 0, k_buf_offset = 0, v_buf_offset = 0, d_buf_offset = 0, m_buf_offset = 0; -+ -+ bool Q_uma = false, K_uma = false, V_uma = false, D_uma = false, M_uma = false; -+ -+ if (ctx->device->uma) { -+ ggml_vk_host_get(ctx->device, q->data, d_Q, q_buf_offset); -+ ggml_vk_host_get(ctx->device, k->data, d_K, q_buf_offset); -+ ggml_vk_host_get(ctx->device, v->data, d_V, q_buf_offset); -+ ggml_vk_host_get(ctx->device, dst->data, d_D, q_buf_offset); -+ Q_uma = d_Q != nullptr; -+ K_uma = d_K != nullptr; -+ V_uma = d_V != nullptr; -+ D_uma = d_D != nullptr; -+ if (mask) { -+ ggml_vk_host_get(ctx->device, mask->data, d_M, q_buf_offset); -+ M_uma = d_M != nullptr; -+ } -+ } -+ -+ -+ ggml_backend_vk_buffer_context * d_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; -+ ggml_backend_vk_buffer_context * q_buf_ctx = (ggml_backend_vk_buffer_context *)q->buffer->context; -+ ggml_backend_vk_buffer_context * k_buf_ctx = (ggml_backend_vk_buffer_context *)k->buffer->context; -+ ggml_backend_vk_buffer_context * v_buf_ctx = (ggml_backend_vk_buffer_context *)v->buffer->context; -+ -+ if (!Q_uma) { -+ d_Q = q_buf_ctx->dev_buffer; -+ q_buf_offset = vk_tensor_offset(q) + q->view_offs; -+ } -+ if (!K_uma) { -+ d_K = k_buf_ctx->dev_buffer; -+ k_buf_offset = vk_tensor_offset(k) + k->view_offs; -+ } -+ if (!V_uma) { -+ d_V = v_buf_ctx->dev_buffer; -+ v_buf_offset = vk_tensor_offset(v) + v->view_offs; -+ } -+ if (!D_uma) { -+ d_D = d_buf_ctx->dev_buffer; -+ d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; -+ } -+ -+ if (!M_uma) { -+ d_M = d_Q; -+ m_buf_offset = q_buf_offset; -+ if (mask) { -+ ggml_backend_vk_buffer_context * m_buf_ctx = (ggml_backend_vk_buffer_context*)mask->buffer->context; -+ d_M = m_buf_ctx->dev_buffer; -+ m_buf_offset = vk_tensor_offset(mask) + mask->view_offs; -+ } -+ } -+ -+ const vk_flash_attn_push_constants pc = { N, KV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3, (uint32_t)neq2, (uint32_t)neq3, (uint32_t)nek2, (uint32_t)nek3, (uint32_t)nev2, (uint32_t)nev3, nem1, (uint32_t)nbq2, (uint32_t)nbq3, (uint32_t)nbk2, (uint32_t)nbk3, (uint32_t)nbv2, (uint32_t)nbv3, nbm1, scale, max_bias, logit_softcap, mask != nullptr, n_head_log2, m0, m1 }; -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, -+ { -+ vk_subbuffer{d_Q, q_buf_offset, VK_WHOLE_SIZE}, -+ vk_subbuffer{d_K, k_buf_offset, VK_WHOLE_SIZE}, -+ vk_subbuffer{d_V, v_buf_offset, VK_WHOLE_SIZE}, -+ vk_subbuffer{d_M, m_buf_offset, VK_WHOLE_SIZE}, -+ vk_subbuffer{d_D, d_buf_offset, VK_WHOLE_SIZE}, -+ }, -+ sizeof(vk_flash_attn_push_constants), &pc, { (uint32_t)neq1, (uint32_t)neq2, (uint32_t)neq3 }); -+} -+ -+static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, ggml_op op) { -+ switch (op) { -+ case GGML_OP_GET_ROWS: -+ GGML_ASSERT(src1->type == GGML_TYPE_I32); -+ if (dst->type == GGML_TYPE_F16) { -+ return ctx->device->pipeline_get_rows[src0->type]; -+ } -+ if (dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_get_rows_f32[src0->type]; -+ } -+ return nullptr; -+ case GGML_OP_ACC: -+ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_acc_f32; -+ } -+ return nullptr; -+ case GGML_OP_ADD: -+ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_add_f32_norepeat : ctx->device->pipeline_add_f32; -+ } -+ if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) { -+ return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_add_f16_f32_f16_norepeat : ctx->device->pipeline_add_f16_f32_f16; -+ } -+ return nullptr; -+ case GGML_OP_MUL: -+ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_mul_f32_norepeat : ctx->device->pipeline_mul_f32; -+ } -+ return nullptr; -+ case GGML_OP_DIV: -+ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ggml_are_same_shape(src0, src1) ? ctx->device->pipeline_div_f32_norepeat : ctx->device->pipeline_div_f32; -+ } -+ return nullptr; -+ case GGML_OP_CONCAT: -+ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_concat_f32; -+ } -+ if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { -+ return ctx->device->pipeline_concat_f16; -+ } -+ if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_I32 && dst->type == GGML_TYPE_I32) { -+ return ctx->device->pipeline_concat_i32; -+ } -+ return nullptr; -+ case GGML_OP_UPSCALE: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_upscale_f32; -+ } -+ return nullptr; -+ case GGML_OP_SCALE: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_scale_f32; -+ } -+ return nullptr; -+ case GGML_OP_SQR: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_sqr_f32; -+ } -+ return nullptr; -+ case GGML_OP_SIN: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_sin_f32; -+ } -+ return nullptr; -+ case GGML_OP_COS: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_cos_f32; -+ } -+ return nullptr; -+ case GGML_OP_CLAMP: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_clamp_f32; -+ } -+ return nullptr; -+ case GGML_OP_PAD: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_pad_f32; -+ } -+ return nullptr; -+ case GGML_OP_REPEAT: -+ if (ggml_type_size(src0->type) == sizeof(float) && ggml_type_size(dst->type) == sizeof(float)) { -+ return ctx->device->pipeline_repeat_f32; -+ } -+ return nullptr; -+ case GGML_OP_CPY: -+ case GGML_OP_CONT: -+ case GGML_OP_DUP: -+ return ggml_vk_get_cpy_pipeline(ctx, src0, dst, dst->type); -+ case GGML_OP_NORM: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_norm_f32; -+ } -+ return nullptr; -+ case GGML_OP_GROUP_NORM: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_group_norm_f32; -+ } -+ return nullptr; -+ case GGML_OP_RMS_NORM: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_rms_norm_f32; -+ } -+ return nullptr; -+ case GGML_OP_UNARY: -+ switch (ggml_get_unary_op(dst)) { -+ case GGML_UNARY_OP_SILU: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_silu_f32; -+ } -+ break; -+ case GGML_UNARY_OP_GELU: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_gelu_f32; -+ } -+ break; -+ case GGML_UNARY_OP_GELU_QUICK: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_gelu_quick_f32; -+ } -+ break; -+ case GGML_UNARY_OP_RELU: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_relu_f32; -+ } -+ break; -+ case GGML_UNARY_OP_TANH: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_tanh_f32; -+ } -+ break; -+ default: -+ break; -+ } -+ return nullptr; -+ case GGML_OP_DIAG_MASK_INF: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_diag_mask_inf_f32; -+ } -+ return nullptr; -+ case GGML_OP_SOFT_MAX: -+ GGML_ASSERT(!src1 || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); -+ -+ if (src0->type == GGML_TYPE_F32 && (src1 == nullptr || src1->type == GGML_TYPE_F32) && dst->type == GGML_TYPE_F32) { -+ return src0->ne[0] > 1024 ? ctx->device->pipeline_soft_max_f32_wg512 : ctx->device->pipeline_soft_max_f32; -+ } -+ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) { -+ return src0->ne[0] > 1024 ? ctx->device->pipeline_soft_max_f32_f16_wg512 : ctx->device->pipeline_soft_max_f32_f16; -+ } -+ return nullptr; -+ case GGML_OP_ROPE: -+ { -+ const int mode = ((const int32_t *) dst->op_params)[2]; -+ const bool is_neox = mode & GGML_ROPE_TYPE_NEOX; -+ -+ if (is_neox) { -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_rope_neox_f32; -+ } -+ if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { -+ return ctx->device->pipeline_rope_neox_f16; -+ } -+ } else { -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_rope_norm_f32; -+ } -+ if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { -+ return ctx->device->pipeline_rope_norm_f16; -+ } -+ } -+ return nullptr; -+ } -+ case GGML_OP_ARGSORT: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_I32) { -+ return ctx->device->pipeline_argsort_f32; -+ } -+ return nullptr; -+ case GGML_OP_SUM_ROWS: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_sum_rows_f32; -+ } -+ return nullptr; -+ case GGML_OP_IM2COL: -+ if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_im2col_f32; -+ } -+ if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) { -+ return ctx->device->pipeline_im2col_f32_f16; -+ } -+ return nullptr; -+ case GGML_OP_TIMESTEP_EMBEDDING: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_timestep_embedding_f32; -+ } -+ return nullptr; -+ case GGML_OP_POOL_2D: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_pool2d_f32; -+ } -+ return nullptr; -+ case GGML_OP_RWKV_WKV6: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_rwkv_wkv6_f32; -+ } -+ return nullptr; -+ case GGML_OP_LEAKY_RELU: -+ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { -+ return ctx->device->pipeline_leaky_relu_f32; -+ } -+ return nullptr; -+ default: -+ return nullptr; -+ } -+ -+ GGML_UNUSED(src2); -+} -+ -+static bool ggml_vk_op_supports_incontiguous(ggml_op op) { -+ switch (op) { -+ case GGML_OP_CPY: -+ case GGML_OP_GET_ROWS: -+ case GGML_OP_ADD: -+ case GGML_OP_MUL: -+ case GGML_OP_DIV: -+ case GGML_OP_CONCAT: -+ case GGML_OP_UPSCALE: -+ case GGML_OP_SQR: -+ case GGML_OP_SIN: -+ case GGML_OP_COS: -+ case GGML_OP_CLAMP: -+ case GGML_OP_PAD: -+ case GGML_OP_REPEAT: -+ return true; -+ default: -+ return false; -+ } -+} -+ -+static uint32_t get_misalign_bytes(ggml_backend_vk_context * ctx, const ggml_tensor * t) -+{ -+ return ((vk_tensor_offset(t) + t->view_offs) & (ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1));; -+} -+ -+template void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, T &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { -+ GGML_UNUSED(p); -+ GGML_UNUSED(src0); -+ GGML_UNUSED(src1); -+ GGML_UNUSED(src2); -+ GGML_UNUSED(dst); -+ static_assert(!std::is_const::value, "unexpected type"); -+ GGML_ASSERT(!src0 || get_misalign_bytes(ctx, src0) == 0); -+ GGML_ASSERT(!src1 || get_misalign_bytes(ctx, src1) == 0); -+ GGML_ASSERT(!src2 || get_misalign_bytes(ctx, src2) == 0); -+ GGML_ASSERT(!dst || get_misalign_bytes(ctx, dst) == 0); -+} -+ -+template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_unary_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { -+ const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); -+ const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); -+ -+ p.misalign_offsets = (a_offset << 16) | d_offset; -+ -+ GGML_UNUSED(src1); -+ GGML_UNUSED(src2); -+} -+ -+template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_binary_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { -+ const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); -+ const uint32_t b_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); -+ const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); -+ -+ GGML_ASSERT(dst->op != GGML_OP_GET_ROWS || (a_offset == 0 && b_offset == 0 && d_offset == 0)); -+ -+ p.misalign_offsets = (a_offset << 16) | (b_offset << 8) | d_offset; -+ -+ GGML_UNUSED(src2); -+} -+ -+template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_upscale_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) { -+ const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); -+ const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); -+ -+ p.a_offset = a_offset; -+ p.d_offset = d_offset; -+ -+ GGML_UNUSED(src1); -+ GGML_UNUSED(src2); -+} -+ -+template -+static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, ggml_op op, PC&& pc, bool dryrun = false) { -+ VK_LOG_DEBUG("ggml_vk_op_f32((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; -+ if (src1 != nullptr) { -+ std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; -+ } -+ if (src2 != nullptr) { -+ std::cerr << "), (" << src2 << ", name=" << src2->name << ", type=" << src2->type << ", ne0=" << src2->ne[0] << ", ne1=" << src2->ne[1] << ", ne2=" << src2->ne[2] << ", ne3=" << src2->ne[3] << ", nb0=" << src2->nb[0] << ", nb1=" << src2->nb[1] << ", nb2=" << src2->nb[2] << ", nb3=" << src2->nb[3]; -+ } -+ std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; -+ std::cerr << "), " << ggml_op_name(op) << ", " << (dryrun ? "dryrun" : "") << ")"); -+ GGML_ASSERT(op == GGML_OP_GET_ROWS || (!ggml_is_quantized(src0->type) && (src1 == nullptr || !ggml_is_quantized(src1->type)))); // NOLINT -+ GGML_ASSERT(ggml_vk_op_supports_incontiguous(op) || ggml_vk_dim01_contiguous(src0)); // NOLINT -+ GGML_ASSERT(dst->buffer != nullptr); -+ const uint64_t ne00 = src0->ne[0]; -+ const uint64_t ne01 = src0->ne[1]; -+ const uint64_t ne02 = src0->ne[2]; -+ const uint64_t ne03 = src0->ne[3]; -+ const uint64_t ne0 = ne00 * ne01; -+ -+ const bool use_src1 = src1 != nullptr; -+ const uint64_t ne10 = use_src1 ? src1->ne[0] : 0; -+ const uint64_t ne11 = use_src1 ? src1->ne[1] : 0; -+ const uint64_t ne12 = use_src1 ? src1->ne[2] : 0; -+ const uint64_t ne13 = use_src1 ? src1->ne[3] : 0; -+ const uint64_t ne1 = ne10 * ne11; -+ // const uint64_t nb10 = use_src1 ? src1->nb[0] : 0; -+ -+ const bool use_src2 = src2 != nullptr; -+ const uint64_t ne20 = use_src2 ? src2->ne[0] : 0; -+ const uint64_t ne21 = use_src2 ? src2->ne[1] : 0; -+ const uint64_t ne22 = use_src2 ? src2->ne[2] : 0; -+ const uint64_t ne23 = use_src2 ? src2->ne[3] : 0; -+ const uint64_t ne2 = ne20 * ne21; -+ -+ const uint64_t ned0 = dst->ne[0]; -+ const uint64_t ned1 = dst->ne[1]; -+ const uint64_t ned2 = dst->ne[2]; -+ const uint64_t ned3 = dst->ne[3]; -+ const uint64_t ned = ned0 * ned1; -+ -+ init_pushconst_fastdiv(pc); -+ -+ vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, src2, dst, op); -+ -+ if (pipeline == nullptr) { -+ std::cerr << "ggml_vulkan: Error: Missing op: " << ggml_op_name(op) << " for " << ggml_type_name(src0->type); -+ if (src1 != nullptr) { -+ std::cerr << " and " << ggml_type_name(src1->type); -+ } -+ std::cerr << " to " << ggml_type_name(dst->type) << std::endl; -+ GGML_ABORT("fatal error"); -+ } -+ -+ if (dryrun) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); -+ return; -+ } -+ -+ const bool op_supports_incontiguous = ggml_vk_op_supports_incontiguous(op); -+ -+ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; -+ ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; -+ ggml_backend_vk_buffer_context * src1_buf_ctx = use_src1 ? (ggml_backend_vk_buffer_context *)src1->buffer->context : nullptr; -+ ggml_backend_vk_buffer_context * src2_buf_ctx = use_src2 ? (ggml_backend_vk_buffer_context *)src2->buffer->context : nullptr; -+ -+ vk_buffer d_X = nullptr; -+ size_t x_buf_offset = 0; -+ vk_buffer d_Y = nullptr; -+ size_t y_buf_offset = 0; -+ vk_buffer d_Z = nullptr; -+ size_t z_buf_offset = 0; -+ -+ bool src0_uma = false; -+ bool src1_uma = false; -+ bool src2_uma = false; -+ -+ if (ctx->device->uma) { -+ ggml_vk_host_get(ctx->device, src0->data, d_X, x_buf_offset); -+ src0_uma = d_X != nullptr; -+ if (use_src1) { -+ ggml_vk_host_get(ctx->device, src1->data, d_Y, y_buf_offset); -+ src1_uma = d_Y != nullptr; -+ } -+ if (use_src2) { -+ ggml_vk_host_get(ctx->device, src2->data, d_Z, z_buf_offset); -+ src2_uma = d_Z != nullptr; -+ } -+ } -+ -+ uint64_t x_sz = ggml_type_size(src0->type)/ggml_blck_size(src0->type) * ne0; -+ uint64_t y_sz = use_src1 ? ggml_type_size(src1->type) * ne1 : 0; -+ uint64_t z_sz = use_src2 ? ggml_type_size(src2->type) * ne2 : 0; -+ uint64_t d_sz = ggml_type_size(dst->type) * ned; -+ -+ vk_buffer d_D = dst_buf_ctx->dev_buffer; -+ -+ // Workaround for tiny tensor inputs on ROPE -+ if (op == GGML_OP_ROPE && use_src1 && y_sz > d_D->size) { -+ y_sz = VK_WHOLE_SIZE; -+ } -+ -+ GGML_ASSERT(d_D != nullptr); -+ uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; -+ if(!src0_uma) { -+ d_X = src0_buf_ctx->dev_buffer; -+ x_buf_offset = vk_tensor_offset(src0) + src0->view_offs; -+ GGML_ASSERT(d_X != nullptr); -+ } -+ if (use_src1 && !src1_uma) { -+ d_Y = src1_buf_ctx->dev_buffer; -+ y_buf_offset = vk_tensor_offset(src1) + src1->view_offs; -+ GGML_ASSERT(d_Y != nullptr); -+ } -+ if (use_src2 && !src2_uma) { -+ d_Z = src2_buf_ctx->dev_buffer; -+ z_buf_offset = vk_tensor_offset(src2) + src2->view_offs; -+ GGML_ASSERT(d_Z != nullptr); -+ } -+ // Compute misalignment offset for descriptors and store it in in push constants, then align the descriptor offsets. -+ init_pushconst_tensor_offsets(ctx, pc, src0, src1, src2, dst); -+ x_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); -+ y_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); -+ z_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); -+ d_buf_offset &= ~(ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); -+ -+ if (op_supports_incontiguous) { -+ x_sz = ggml_nbytes(src0); -+ y_sz = use_src1 ? ggml_nbytes(src1) : 0; -+ z_sz = use_src2 ? ggml_nbytes(src2) : 0; -+ d_sz = ggml_nbytes(dst); -+ -+ if (x_buf_offset + x_sz >= d_X->size) { -+ x_sz = VK_WHOLE_SIZE; -+ } -+ if (use_src1 && y_buf_offset + y_sz >= d_Y->size) { -+ y_sz = VK_WHOLE_SIZE; -+ } -+ if (use_src2 && z_buf_offset + z_sz >= d_Z->size) { -+ z_sz = VK_WHOLE_SIZE; -+ } -+ if (d_buf_offset + d_sz >= d_D->size) { -+ d_sz = VK_WHOLE_SIZE; -+ } -+ } -+ -+ std::array elements; -+ -+ // Single call if dimension 2 is contiguous -+ GGML_ASSERT(op_supports_incontiguous || (ggml_is_contiguous(src0) && (src1 == nullptr || ggml_is_contiguous(src1)))); -+ -+ switch (op) { -+ case GGML_OP_NORM: -+ case GGML_OP_RMS_NORM: -+ case GGML_OP_SOFT_MAX: -+ case GGML_OP_SUM_ROWS: -+ { -+ const uint32_t nr = ggml_nrows(src0); -+ if (nr > 262144) { -+ elements = { 512, 512, CEIL_DIV(nr, 262144) }; -+ } else if (nr > 512) { -+ elements = { 512, CEIL_DIV(nr, 512), 1 }; -+ } else { -+ elements = { nr, 1, 1 }; -+ } -+ } break; -+ case GGML_OP_GROUP_NORM: -+ { -+ const uint32_t num_groups = dst->op_params[0]; -+ elements = { num_groups * (uint32_t)src0->ne[3], 1, 1 }; -+ } break; -+ case GGML_OP_DIAG_MASK_INF: -+ case GGML_OP_ROPE: -+ elements = { (uint32_t)ggml_nrows(src0), (uint32_t)ne00, 1 }; -+ break; -+ case GGML_OP_GET_ROWS: -+ elements = { (uint32_t)ne00, (uint32_t)ne10, (uint32_t)(ne11 * ne12) }; -+ break; -+ case GGML_OP_ARGSORT: -+ elements = { (uint32_t)ne00, (uint32_t)ggml_nrows(src0), 1 }; -+ break; -+ case GGML_OP_IM2COL: -+ { -+ const bool is_2D = dst->op_params[6] == 1; -+ -+ const uint32_t IC = src1->ne[is_2D ? 2 : 1]; -+ -+ const uint32_t KH = is_2D ? src0->ne[1] : 1; -+ const uint32_t KW = src0->ne[0]; -+ -+ const uint32_t OH = is_2D ? dst->ne[2] : 1; -+ const uint32_t OW = dst->ne[1]; -+ -+ const uint32_t batch = src1->ne[is_2D ? 3 : 2]; -+ -+ elements = { OW * KW * KH, OH, batch * IC }; -+ } break; -+ case GGML_OP_TIMESTEP_EMBEDDING: -+ { -+ const uint32_t dim = dst->op_params[0]; -+ uint32_t half_ceil = (dim + 1) / 2; -+ elements = { half_ceil, (uint32_t)src0->ne[0], 1 }; -+ } break; -+ case GGML_OP_POOL_2D: -+ { -+ const uint32_t N = dst->ne[3]; -+ const uint32_t OC = dst->ne[2]; -+ const uint32_t OH = dst->ne[1]; -+ const uint32_t OW = dst->ne[0]; -+ elements = { N * OC * OH * OW, 1, 1}; -+ } break; -+ case GGML_OP_ADD: -+ case GGML_OP_DIV: -+ case GGML_OP_MUL: -+ case GGML_OP_SCALE: -+ case GGML_OP_SQR: -+ case GGML_OP_SIN: -+ case GGML_OP_COS: -+ case GGML_OP_CLAMP: -+ case GGML_OP_PAD: -+ case GGML_OP_REPEAT: -+ case GGML_OP_CPY: -+ case GGML_OP_CONCAT: -+ case GGML_OP_UPSCALE: -+ case GGML_OP_UNARY: -+ { -+ const uint32_t ne = ggml_nelements(dst); -+ if (ne > 262144) { -+ elements = { 512, 512, CEIL_DIV(ne, 262144) }; -+ } else if (ne > 512) { -+ elements = { 512, CEIL_DIV(ne, 512), 1 }; -+ } else { -+ elements = { ne, 1, 1 }; -+ } -+ } break; -+ default: -+ elements = { (uint32_t)ggml_nelements(src0), 1, 1 }; -+ break; -+ } -+ -+ if (!op_supports_incontiguous) { -+ if (x_sz != VK_WHOLE_SIZE) { -+ x_sz *= ne02 * ne03; -+ } -+ if (use_src1 && y_sz != VK_WHOLE_SIZE) { -+ y_sz *= ne12 * ne13; -+ } -+ if (use_src2 && z_sz != VK_WHOLE_SIZE) { -+ z_sz *= ne22 * ne23; -+ } -+ if (d_sz != VK_WHOLE_SIZE) { -+ d_sz *= ned2 * ned3; -+ } -+ } -+ -+ if (op == GGML_OP_SOFT_MAX) { -+ // Empty src1 is possible in soft_max, but the shader needs a buffer -+ vk_subbuffer subbuf_y; -+ if (use_src1) { -+ subbuf_y = { d_Y, y_buf_offset, y_sz }; -+ } else { -+ subbuf_y = { d_X, 0, x_sz }; -+ } -+ -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, subbuf_y, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); -+ } else if (op == GGML_OP_ROPE) { -+ // Empty src2 is possible in rope, but the shader needs a buffer -+ vk_subbuffer subbuf_z; -+ if (use_src2) { -+ subbuf_z = { d_Z, z_buf_offset, z_sz }; -+ } else { -+ subbuf_z = { d_X, 0, x_sz }; -+ } -+ -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, subbuf_z, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); -+ } else if (op == GGML_OP_IM2COL) { -+ // im2col uses only src1 and dst buffers -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); -+ } else if (use_src2) { -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_Z, z_buf_offset, z_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); -+ } else if (use_src1) { -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_Y, y_buf_offset, y_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); -+ } else { -+ ggml_vk_sync_buffers(subctx); -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { vk_subbuffer{ d_X, x_buf_offset, x_sz }, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, sizeof(PC), &pc, elements); -+ } -+} -+ -+static void ggml_vk_get_rows(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t src1_type_size = ggml_type_size(src1->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_GET_ROWS, { -+ (uint32_t)ggml_nelements(src0), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_acc(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t src1_type_size = ggml_type_size(src1->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ int nb1 = dst->op_params[0] / 4; // 4 bytes of float32 -+ int nb2 = dst->op_params[1] / 4; // 4 bytes of float32 -+ // int nb3 = dst->op_params[2] / 4; // 4 bytes of float32 - unused -+ int offset = dst->op_params[3] / 4; // offset in bytes -+ -+ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_ACC, { -+ (uint32_t)ggml_nelements(src0), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)nb1, (uint32_t)nb2, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t)nb1, (uint32_t)nb2, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, offset, -+ }, dryrun); -+} -+ -+static void ggml_vk_add(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t src1_type_size = ggml_type_size(src1->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_ADD, { -+ (uint32_t)ggml_nelements(src0), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_mul(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t src1_type_size = ggml_type_size(src1->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_MUL, { -+ (uint32_t)ggml_nelements(src0), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_div(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t src1_type_size = ggml_type_size(src1->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_DIV, { -+ (uint32_t)ggml_nelements(src0), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_op_f32_rwkv6(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, const vk_op_rwkv_wkv6_push_constants&& pc, bool dryrun = false) { -+ const ggml_tensor * k = dst->src[0]; -+ const ggml_tensor * v = dst->src[1]; -+ const ggml_tensor * r = dst->src[2]; -+ const ggml_tensor * tf = dst->src[3]; -+ const ggml_tensor * td = dst->src[4]; -+ const ggml_tensor * state = dst->src[5]; -+ -+ GGML_ASSERT(!ggml_is_quantized(k->type)); -+ GGML_ASSERT(!ggml_is_quantized(v->type)); -+ GGML_ASSERT(!ggml_is_quantized(r->type)); -+ GGML_ASSERT(!ggml_is_quantized(tf->type)); -+ GGML_ASSERT(!ggml_is_quantized(td->type)); -+ GGML_ASSERT(!ggml_is_quantized(state->type)); -+ GGML_ASSERT(dst->buffer != nullptr); -+ -+ vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, k, v, r, dst, GGML_OP_RWKV_WKV6); -+ GGML_ASSERT(pipeline != nullptr); -+ -+ if (dryrun) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); -+ return; -+ } -+ -+ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; -+ ggml_backend_vk_buffer_context * k_buf_ctx = (ggml_backend_vk_buffer_context *)k->buffer->context; -+ ggml_backend_vk_buffer_context * v_buf_ctx = (ggml_backend_vk_buffer_context *)v->buffer->context; -+ ggml_backend_vk_buffer_context * r_buf_ctx = (ggml_backend_vk_buffer_context *)r->buffer->context; -+ ggml_backend_vk_buffer_context * tf_buf_ctx = (ggml_backend_vk_buffer_context *)tf->buffer->context; -+ ggml_backend_vk_buffer_context * td_buf_ctx = (ggml_backend_vk_buffer_context *)td->buffer->context; -+ ggml_backend_vk_buffer_context * state_buf_ctx = (ggml_backend_vk_buffer_context *)state->buffer->context; -+ -+ ggml_vk_sync_buffers(subctx); -+ -+ vk_buffer d_D = nullptr, d_K = nullptr, d_V = nullptr, d_R = nullptr, d_TF = nullptr, d_TD = nullptr, d_State = nullptr; -+ size_t k_offset = 0, v_offset = 0, r_offset = 0, tf_offset = 0, td_offset = 0, state_offset = 0, dst_offset = 0; -+ bool K_uma = false, V_uma = false, R_uma = false, TF_uma = false, TD_uma = false, STATE_uma = false, DST_uma = false; -+ -+ if (ctx->device->uma) { -+ ggml_vk_host_get(ctx->device, k->data, d_K, k_offset); -+ ggml_vk_host_get(ctx->device, v->data, d_V, v_offset); -+ ggml_vk_host_get(ctx->device, r->data, d_R, r_offset); -+ ggml_vk_host_get(ctx->device, tf->data, d_TF, tf_offset); -+ ggml_vk_host_get(ctx->device, td->data, d_TD, td_offset); -+ ggml_vk_host_get(ctx->device, state->data, d_State, state_offset); -+ ggml_vk_host_get(ctx->device, dst->data, d_D, dst_offset); -+ -+ K_uma = d_K != nullptr; -+ V_uma = d_V != nullptr; -+ R_uma = d_R != nullptr; -+ TF_uma = d_TF != nullptr; -+ TD_uma = d_TD != nullptr; -+ STATE_uma = d_State != nullptr; -+ DST_uma = d_D != nullptr; -+ } -+ -+ if (!K_uma) { -+ d_K = k_buf_ctx->dev_buffer; -+ k_offset = vk_tensor_offset(k) + k->view_offs; -+ } -+ if (!V_uma) { -+ d_V = v_buf_ctx->dev_buffer; -+ v_offset = vk_tensor_offset(v) + v->view_offs; -+ } -+ if (!R_uma) { -+ d_R = r_buf_ctx->dev_buffer; -+ r_offset = vk_tensor_offset(r) + r->view_offs; -+ } -+ if (!TF_uma) { -+ d_TF = tf_buf_ctx->dev_buffer; -+ tf_offset = vk_tensor_offset(tf) + tf->view_offs; -+ } -+ if (!TD_uma) { -+ d_TD = td_buf_ctx->dev_buffer; -+ td_offset = vk_tensor_offset(td) + td->view_offs; -+ } -+ if (!STATE_uma) { -+ d_State = state_buf_ctx->dev_buffer; -+ state_offset = vk_tensor_offset(state) + state->view_offs; -+ } -+ if (!DST_uma) { -+ d_D = dst_buf_ctx->dev_buffer; -+ dst_offset = vk_tensor_offset(dst) + dst->view_offs; -+ } -+ -+ const uint64_t k_size = ggml_nbytes(k); -+ const uint64_t v_size = ggml_nbytes(v); -+ const uint64_t r_size = ggml_nbytes(r); -+ const uint64_t tf_size = ggml_nbytes(tf); -+ const uint64_t td_size = ggml_nbytes(td); -+ const uint64_t state_size = ggml_nbytes(state); -+ const uint64_t dst_size = ggml_nbytes(dst); -+ -+ std::array elements = { -+ (uint32_t)(pc.B * pc.H), -+ 1, -+ 1 -+ }; -+ -+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { -+ vk_subbuffer{ d_K, k_offset, k_size }, -+ vk_subbuffer{ d_V, v_offset, v_size }, -+ vk_subbuffer{ d_R, r_offset, r_size }, -+ vk_subbuffer{ d_TF, tf_offset, tf_size }, -+ vk_subbuffer{ d_TD, td_offset, td_size }, -+ vk_subbuffer{ d_State, state_offset, state_size }, -+ vk_subbuffer{ d_D, dst_offset, dst_size } -+ }, sizeof(vk_op_rwkv_wkv6_push_constants), &pc, elements); -+} -+ -+static void ggml_vk_rwkv_wkv6(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, bool dryrun = false) { -+ const size_t seq_length = dst->src[0]->ne[3]; -+ const size_t n_embed = dst->ne[0]; -+ const size_t n_heads = dst->src[0]->ne[2]; -+ const size_t n_seqs = dst->src[5]->ne[1]; -+ -+ ggml_vk_op_f32_rwkv6( -+ ctx, subctx, dst, -+ { -+ (uint32_t)n_seqs, -+ (uint32_t)seq_length, -+ (uint32_t)n_embed, -+ (uint32_t)n_heads, -+ }, -+ dryrun -+ ); -+} -+ -+static void ggml_vk_concat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ int * op_params = (int *)dst->op_params; -+ -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t src1_type_size = ggml_type_size(src1->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_CONCAT, { -+ (uint32_t)ggml_nelements(dst), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, op_params[0], -+ }, dryrun); -+} -+ -+static void ggml_vk_upscale(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ -+ const float sf0 = (float)dst->ne[0] / src0->ne[0]; -+ const float sf1 = (float)dst->ne[1] / src0->ne[1]; -+ const float sf2 = (float)dst->ne[2] / src0->ne[2]; -+ const float sf3 = (float)dst->ne[3] / src0->ne[3]; -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_UPSCALE, { -+ (uint32_t)ggml_nelements(dst), 0, 0, -+ (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t)dst->ne[0], (uint32_t)dst->ne[1], (uint32_t)dst->ne[2],(uint32_t)dst->ne[3], -+ sf0, sf1, sf2, sf3, -+ }, dryrun); -+} -+ -+static void ggml_vk_scale(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ float * op_params = (float *)dst->op_params; -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SCALE, { -+ (uint32_t)ggml_nelements(src0), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ op_params[0], 0.0f, -+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_sqr(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SQR, { -+ (uint32_t)ggml_nelements(src0), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, -+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_sin(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SIN, { -+ (uint32_t)ggml_nelements(src0), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, -+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_cos(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_COS, { -+ (uint32_t)ggml_nelements(src0), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, -+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_clamp(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ float * op_params = (float *)dst->op_params; -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_CLAMP, { -+ (uint32_t)ggml_nelements(src0), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ op_params[0], op_params[1], -+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_pad(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_PAD, { -+ (uint32_t)ggml_nelements(dst), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, -+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_repeat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_REPEAT, { -+ (uint32_t)ggml_nelements(dst), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, -+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_cpy(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t src0_type_size = ggml_type_size(src0->type); -+ const uint32_t dst_type_size = ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_CPY, { -+ (uint32_t)ggml_nelements(src0), -+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, -+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, -+ 0, -+ 0.0f, 0.0f, -+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -+ }, dryrun); -+} -+ -+static void ggml_vk_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ float * op_params = (float *)dst->op_params; -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_NORM, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0], 0.0f }, dryrun); -+} -+ -+static void ggml_vk_group_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ const int * int_op_params = (const int *)dst->op_params; -+ const float * float_op_params = (const float *)dst->op_params; -+ -+ const uint32_t num_groups = int_op_params[0]; -+ const float eps = float_op_params[1]; -+ const uint32_t group_size = src0->ne[0] * src0->ne[1] * ((src0->ne[2] + num_groups - 1) / num_groups); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_GROUP_NORM, { group_size, 0, eps, 0.0f }, dryrun); -+} -+ -+static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ float * op_params = (float *)dst->op_params; -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_RMS_NORM, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0], 0.0f }, dryrun); -+} -+ -+static void ggml_vk_unary(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_UNARY, { (uint32_t)ggml_nelements(src0), 0, 0.0f, 0.0f }, dryrun); -+} -+ -+static void ggml_vk_diag_mask_inf(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ int32_t * op_params = (int32_t *)dst->op_params; -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_DIAG_MASK_INF, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0] }, dryrun); -+} -+ -+static void ggml_vk_soft_max(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ float * op_params = (float *)dst->op_params; -+ -+ float scale = op_params[0]; -+ float max_bias = op_params[1]; -+ -+ const uint32_t ncols = (uint32_t)src0->ne[0]; -+ const uint32_t nrows_x = (uint32_t)ggml_nrows(src0); -+ const uint32_t nrows_y = (uint32_t)src0->ne[1]; -+ -+ const uint32_t n_head_kv = nrows_x/nrows_y; -+ const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head_kv)); -+ -+ const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); -+ const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_SOFT_MAX, { -+ ncols, -+ src1 != nullptr ? nrows_y : (uint32_t)0, -+ scale, max_bias, -+ m0, m1, -+ n_head_log2, -+ nrows_x, -+ }, dryrun); -+} -+ -+static void ggml_vk_rope(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, bool dryrun = false) { -+ const int n_dims = ((int32_t *) dst->op_params)[1]; -+ // const int mode = ((int32_t *) dst->op_params)[2]; -+ // const int n_ctx = ((int32_t *) dst->op_params)[3]; -+ const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; -+ const float freq_base = ((float *) dst->op_params)[5]; -+ const float freq_scale = ((float *) dst->op_params)[6]; -+ const float ext_factor = ((float *) dst->op_params)[7]; -+ const float attn_factor = ((float *) dst->op_params)[8]; -+ const float beta_fast = ((float *) dst->op_params)[9]; -+ const float beta_slow = ((float *) dst->op_params)[10]; -+ -+ float corr_dims[2]; -+ ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); -+ -+ const float theta_scale = powf(freq_base, -2.0f/n_dims); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, src1, src2, dst, GGML_OP_ROPE, { -+ (uint32_t)src0->ne[0], (uint32_t)n_dims, freq_scale, (uint32_t)src0->ne[1], -+ freq_base, ext_factor, attn_factor, {corr_dims[0], corr_dims[1]}, theta_scale, -+ src2 != nullptr, -+ }, dryrun); -+} -+ -+static void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ int32_t * op_params = (int32_t *)dst->op_params; -+ -+ uint32_t ncols = src0->ne[0]; -+ -+ uint32_t ncols_pad = 1; -+ while (ncols_pad < ncols) { -+ ncols_pad *= 2; -+ } -+ -+ GGML_ASSERT(ncols_pad <= 1024); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_ARGSORT, { -+ ncols, -+ ncols_pad, -+ op_params[0], -+ }, dryrun); -+} -+ -+static void ggml_vk_sum_rows(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_SUM_ROWS, { (uint32_t)src0->ne[0], 0, 0.0f, 0.0f }, dryrun); -+} -+ -+static void ggml_vk_im2col(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool dryrun = false) { -+ const int32_t s0 = dst->op_params[0]; -+ const int32_t s1 = dst->op_params[1]; -+ const int32_t p0 = dst->op_params[2]; -+ const int32_t p1 = dst->op_params[3]; -+ const int32_t d0 = dst->op_params[4]; -+ const int32_t d1 = dst->op_params[5]; -+ -+ const bool is_2D = dst->op_params[6] == 1; -+ -+ const uint32_t IC = src1->ne[is_2D ? 2 : 1]; -+ const uint32_t IH = is_2D ? src1->ne[1] : 1; -+ const uint32_t IW = src1->ne[0]; -+ -+ const uint32_t KH = is_2D ? src0->ne[1] : 1; -+ const uint32_t KW = src0->ne[0]; -+ -+ const uint32_t OH = is_2D ? dst->ne[2] : 1; -+ const uint32_t OW = dst->ne[1]; -+ -+ const uint32_t offset_delta = src1->nb[is_2D ? 2 : 1] / 4; // nb is byte offset, src is type float32 -+ const uint32_t batch_offset = src1->nb[is_2D ? 3 : 2] / 4; // nb is byte offset, src is type float32 -+ -+ const uint32_t pelements = OW * KW * KH; -+ -+ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_IM2COL, { -+ batch_offset, offset_delta, -+ IC, IW, IH, OW, OH, KW, KH, -+ pelements, -+ IC * KH * KW, -+ s0, s1, p0, p1, d0, d1, -+ }, dryrun); -+} -+ -+static void ggml_vk_timestep_embedding(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ const uint32_t dim = dst->op_params[0]; -+ const uint32_t max_period = dst->op_params[1]; -+ const uint32_t nb1 = dst->nb[1] / ggml_type_size(dst->type); -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_TIMESTEP_EMBEDDING, { -+ nb1, dim, max_period, -+ }, dryrun); -+} -+ -+static void ggml_vk_pool_2d(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ uint32_t op = static_cast(dst->op_params[0]); -+ const int32_t k1 = dst->op_params[1]; -+ const int32_t k0 = dst->op_params[2]; -+ const int32_t s1 = dst->op_params[3]; -+ const int32_t s0 = dst->op_params[4]; -+ const int32_t p1 = dst->op_params[5]; -+ const int32_t p0 = dst->op_params[6]; -+ -+ const uint32_t IH = src0->ne[1]; -+ const uint32_t IW = src0->ne[0]; -+ -+ const uint32_t N = dst->ne[3]; -+ -+ const uint32_t OC = dst->ne[2]; -+ const uint32_t OH = dst->ne[1]; -+ const uint32_t OW = dst->ne[0]; -+ -+ const uint32_t parallel_elements = N * OC * OH * OW; -+ -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_POOL_2D, { -+ IW, IH, OW, OH, OC, -+ parallel_elements, -+ op, -+ k0, k1, s0, s1, p0, p1, -+ }, dryrun); -+} -+ -+static void ggml_vk_leaky_relu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) { -+ const float * op_params = (const float *)dst->op_params; -+ ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_LEAKY_RELU, { (uint32_t)ggml_nelements(src0), 0, op_params[0], 0.0f }, dryrun); -+} -+ -+#ifdef GGML_VULKAN_RUN_TESTS -+static void ggml_vk_print_matrix_area(const void * data, ggml_type type, int ne0, int ne1, int i0, int i1, int i2) { -+ if (type != GGML_TYPE_F32 && type != GGML_TYPE_F16) { -+ return; -+ } -+ i0 = std::max(i0, 5); -+ i1 = std::max(i1, 5); -+ i2 = std::max(i2, 0); -+ fprintf(stderr, " "); -+ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { -+ fprintf(stderr, "%7d ", idx1); -+ } -+ fprintf(stderr, "\n"); -+ for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { -+ fprintf(stderr, "%7d: ", idx0); -+ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { -+ if (idx0 >= 0 && idx0 < ne0 && idx1 >= 0 && idx1 < ne1) { -+ float val; -+ if (type == GGML_TYPE_F32) { -+ val = *((const float *) data + i2*ne1*ne0 + idx1*ne0 + idx0); -+ } else if (type == GGML_TYPE_F16) { -+ val = ggml_fp16_to_fp32(*((const ggml_fp16_t *) data + i2*ne1*ne0 + idx1*ne0 + idx0)); -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ fprintf(stderr, "% 7.2f ", val); -+ } else { -+ fprintf(stderr, " "); -+ } -+ } -+ fprintf(stderr, "\n"); -+ } -+} -+ -+template -+static void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, int split_k, int shader_size) { -+ VK_LOG_DEBUG("ggml_vk_test_matmul(" << m << ", " << n << ", " << k << ", " << batch << ", " << num_it << ", " << split_k << ", " << shader_size << ")"); -+ const size_t x_ne = m * k * batch; -+ const size_t y_ne = k * n * batch; -+ const size_t d_ne = m * n * batch; -+ -+ vk_pipeline p; -+ std::string shname; -+ if (shader_size == 0) { -+ if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32->a_s; -+ shname = "F32_ALIGNED_S"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32_f16->a_s; -+ shname = "F32_F16_ALIGNED_S"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_s; -+ shname = "F16_F32_ALIGNED_S"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16.f32acc->a_s; -+ shname = "F16_ALIGNED_S"; -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ } else if (shader_size == 1) { -+ if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32->a_m; -+ shname = "F32_ALIGNED_M"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32_f16->a_m; -+ shname = "F32_F16_ALIGNED_M"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_m; -+ shname = "F16_F32_ALIGNED_M"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16.f32acc->a_m; -+ shname = "F16_ALIGNED_M"; -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ } else if (shader_size == 2) { -+ if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32->a_l; -+ shname = "F32_ALIGNED_L"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32_f16->a_l; -+ shname = "F32_F16_ALIGNED_L"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_l; -+ shname = "F16_F32_ALIGNED_L"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16.f32acc->a_l; -+ shname = "F16_ALIGNED_L"; -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ } else { -+ GGML_ASSERT(0); -+ } -+ -+ const size_t kpad = ggml_vk_align_size(k, p->align); -+ -+ if (k != kpad) { -+ if (shader_size == 0) { -+ if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32->s; -+ shname = "F32_S"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32_f16->s; -+ shname = "F32_F16_S"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16_f32.f32acc->s; -+ shname = "F16_F32_S"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16.f32acc->s; -+ shname = "F16_S"; -+ } -+ } else if (shader_size == 1) { -+ if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32->m; -+ shname = "F32_M"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32_f16->m; -+ shname = "F32_F16_M"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16_f32.f32acc->m; -+ shname = "F16_F32_M"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16.f32acc->m; -+ shname = "F16_M"; -+ } -+ } else if (shader_size == 2) { -+ if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32->l; -+ shname = "F32_L"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f32_f16->l; -+ shname = "F32_F16_L"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16_f32.f32acc->l; -+ shname = "F16_F32_L"; -+ } else if (std::is_same() && std::is_same()) { -+ p = ctx->device->pipeline_matmul_f16.f32acc->l; -+ shname = "F16_L"; -+ } -+ } -+ } -+ -+ ggml_pipeline_request_descriptor_sets(ctx->device, p, num_it); -+ if (split_k > 1) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, num_it); -+ -+ if (ctx->prealloc_split_k == nullptr || ctx->prealloc_split_k->size < sizeof(float) * d_ne * split_k) { -+ // Resize buffer -+ if (ctx->prealloc_split_k != nullptr) { -+ ggml_vk_destroy_buffer(ctx->prealloc_split_k); -+ } -+ ctx->prealloc_split_k = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne * split_k, vk::MemoryPropertyFlagBits::eDeviceLocal); -+ } -+ } -+ -+ ggml_pipeline_allocate_descriptor_sets(ctx->device); -+ -+ vk_buffer d_X = ggml_vk_create_buffer_check(ctx->device, sizeof(X_TYPE) * x_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); -+ vk_buffer d_Y = ggml_vk_create_buffer_check(ctx->device, sizeof(Y_TYPE) * y_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); -+ vk_buffer d_D = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne, vk::MemoryPropertyFlagBits::eDeviceLocal); -+ -+ X_TYPE* x = (X_TYPE *) malloc(sizeof(X_TYPE) * x_ne); -+ Y_TYPE* y = (Y_TYPE *) malloc(sizeof(Y_TYPE) * y_ne); -+ float* d = (float *) malloc(sizeof(float) * d_ne); -+ -+ for (size_t i = 0; i < x_ne; i++) { -+ if (std::is_same()) { -+ x[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; -+ // x[i] = 1.0f; -+ // x[i] = i + 1; -+ // x[i] = (i % k == i / k) ? 1.0f : 0.0f; -+ } else if (std::is_same()) { -+ x[i] = ggml_fp32_to_fp16((rand() / (float)RAND_MAX) * 2.0f - 1.0f); -+ // x[i] = ggml_fp32_to_fp16(1.0f); -+ // x[i] = ggml_fp32_to_fp16(i + 1); -+ // x[i] = ggml_fp32_to_fp16((i % k == i / k) ? 1.0f : 0.0f); -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ } -+ for (size_t i = 0; i < y_ne; i++) { -+ if (std::is_same()) { -+ y[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; -+ // y[i] = (i % k == i / k) ? 1.0f : 0.0f; -+ // y[i] = i + 1; -+ } else if (std::is_same()) { -+ y[i] = ggml_fp32_to_fp16((rand() / (float)RAND_MAX) * 2.0f - 1.0f); -+ // y[i] = ggml_fp32_to_fp16((i % k == i / k) ? 1.0f : 0.0f); -+ // y[i] = ggml_fp32_to_fp16(i + 1); -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ } -+ -+ ggml_vk_buffer_write(d_X, 0, x, sizeof(X_TYPE) * k * m * batch); -+ ggml_vk_buffer_write(d_Y, 0, y, sizeof(Y_TYPE) * k * n * batch); -+ -+ vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); -+ ggml_vk_ctx_begin(ctx->device, subctx); -+ for (size_t i = 0; i < num_it; i++) { -+ ggml_vk_matmul( -+ ctx, subctx, p, ggml_vk_subbuffer(d_X), ggml_vk_subbuffer(d_Y), ggml_vk_subbuffer(d_D), ggml_vk_subbuffer(ctx->prealloc_split_k), -+ m, n, k, -+ k, k, m, k*m, k*n, m*n, -+ split_k, batch, batch, batch, 1, 1 -+ ); -+ } -+ ggml_vk_ctx_end(subctx); -+ -+ auto begin = std::chrono::high_resolution_clock::now(); -+ ggml_vk_submit(subctx, ctx->fence); -+ VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_matmul waitForFences"); -+ ctx->device->device.resetFences({ ctx->fence }); -+ -+ auto end = std::chrono::high_resolution_clock::now(); -+ double time = std::chrono::duration_cast(end-begin).count() / 1000.0; -+ -+ // copy dst to host -+ ggml_vk_buffer_read(d_D, 0, d, sizeof(float) * d_ne); -+ -+ float * d_chk = (float *) malloc(sizeof(float) * d_ne); -+ -+ ggml_init_params iparams = { -+ /*.mem_size =*/ 1024*1024*1024, -+ /*.mem_buffer =*/ NULL, -+ /*.no_alloc =*/ true, -+ }; -+ -+ ggml_context * ggml_ctx = ggml_init(iparams); -+ -+ ggml_type src0_type; -+ ggml_type src1_type; -+ -+ if (std::is_same()) { -+ src0_type = GGML_TYPE_F32; -+ } else if (std::is_same()) { -+ src0_type = GGML_TYPE_F16; -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ if (std::is_same()) { -+ src1_type = GGML_TYPE_F32; -+ } else if (std::is_same()) { -+ src1_type = GGML_TYPE_F16; -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ -+ ggml_tensor * src0_ggml = ggml_new_tensor_3d(ggml_ctx, src0_type, k, m, batch); -+ ggml_tensor * src1_ggml = ggml_new_tensor_3d(ggml_ctx, src1_type, k, n, batch); -+ ggml_tensor * tensor_ggml = ggml_mul_mat(ggml_ctx, src0_ggml, src1_ggml); -+ -+ src0_ggml->data = x; -+ src1_ggml->data = y; -+ tensor_ggml->data = d_chk; -+ -+ ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); -+ ggml_build_forward_expand(cgraph, tensor_ggml); -+ -+ ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 1); -+ -+ ggml_free(ggml_ctx); -+ -+ double avg_err = 0.0; -+ int first_err_n = -1; -+ int first_err_m = -1; -+ int first_err_b = -1; -+ -+ for (size_t i = 0; i < m*n*batch; i++) { -+ double err = std::fabs(d[i] - d_chk[i]); -+ avg_err += err; -+ -+ if ((err > 0.05f || std::isnan(err)) && first_err_n == -1) { -+ first_err_b = i / (m * n); -+ first_err_n = (i % (m * n)) / m; -+ first_err_m = (i % (m * n)) % m; -+ } -+ } -+ -+ avg_err /= m * n; -+ -+ double tflops = 2.0*m*n*k*batch*num_it / (time / 1000.0) / (1000.0*1000.0*1000.0*1000.0); -+ -+ std::cerr << "TEST " << shname << " m=" << m << " n=" << n << " k=" << k << " batch=" << batch << " split_k=" << split_k << " matmul " << time / num_it << "ms " << tflops << " TFLOPS avg_err=" << avg_err << std::endl; -+ -+ if (avg_err > 0.1 || std::isnan(avg_err)) { -+ std::cerr << "m = " << first_err_m << " n = " << first_err_n << " b = " << first_err_b << std::endl; -+ std::cerr << "Actual result: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(d, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ std::cerr << "Expected result: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(d_chk, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ -+ if (split_k > 1) { -+ float * split_k_buf = (float *) malloc(sizeof(float) * d_ne * split_k); -+ ggml_vk_buffer_read(ctx->prealloc_split_k, 0, split_k_buf, sizeof(float) * d_ne * split_k); -+ -+ std::cerr << "d_buf0: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(split_k_buf, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ -+ std::cerr << "d_buf1: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(split_k_buf + d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ -+ std::cerr << "d_buf2: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(split_k_buf + 2 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ -+ std::cerr << "d_buf3: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(split_k_buf + 3 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ -+ free(split_k_buf); -+ } -+ } -+ -+ free(d_chk); -+ -+ ggml_vk_queue_cleanup(ctx->device, ctx->device->transfer_queue); -+ ggml_vk_queue_cleanup(ctx->device, ctx->device->compute_queue); -+ -+ ggml_vk_destroy_buffer(d_X); -+ ggml_vk_destroy_buffer(d_Y); -+ ggml_vk_destroy_buffer(d_D); -+ -+ ggml_pipeline_cleanup(p); -+ ggml_pipeline_cleanup(ctx->device->pipeline_matmul_split_k_reduce); -+ -+ free(x); -+ free(y); -+ free(d); -+} -+ -+static void ggml_vk_print_tensor_area(const ggml_tensor * tensor, int i0, int i1, int i2, int i3) { -+ if (tensor->type != GGML_TYPE_F32 && tensor->type != GGML_TYPE_F16) { -+ return; -+ } -+ i0 = std::max(i0, 5); -+ i1 = std::max(i1, 5); -+ i2 = std::max(i2, 0); -+ i3 = std::max(i3, 0); -+ fprintf(stderr, " "); -+ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { -+ fprintf(stderr, "%7d ", idx1); -+ } -+ fprintf(stderr, "\n"); -+ for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { -+ fprintf(stderr, "%7d: ", idx0); -+ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { -+ if (idx0 >= 0 && idx0 < tensor->ne[0] && idx1 >= 0 && idx1 < tensor->ne[1] && i2 >= 0 && i2 < tensor->ne[2] && i3 >= 0 && i3 < tensor->ne[3]) { -+ float val; -+ if (tensor->type == GGML_TYPE_F32) { -+ val = *(float *) ((char *) tensor->data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); -+ } else if (tensor->type == GGML_TYPE_F16) { -+ val = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) tensor->data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0])); -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ fprintf(stderr, "% 7.2f ", val); -+ } else { -+ fprintf(stderr, " "); -+ } -+ } -+ fprintf(stderr, "\n"); -+ } -+} -+ -+static void ggml_vk_quantize_data(const float * from, void * to, size_t ne, ggml_type quant) { -+ ggml_quantize_chunk(quant, from, to, 0, 1, ne, nullptr); -+} -+ -+static void ggml_vk_dequantize_data(const void * from, float * to, size_t ne, ggml_type quant) { -+ if (quant == GGML_TYPE_F32) { -+ memcpy(to, from, sizeof(float) * ne); -+ return; -+ } -+ -+ const auto * tt = ggml_get_type_traits(quant); -+ -+ ggml_to_float_t dequant_fn = tt->to_float; -+ -+ dequant_fn(from, to, ne); -+} -+ -+static void ggml_vk_test_dequant(ggml_backend_vk_context * ctx, size_t ne, ggml_type quant) { -+ VK_LOG_DEBUG("ggml_vk_test_dequant(" << ne << ")"); -+ const size_t x_sz = sizeof(float) * ne; -+ const size_t x_sz_f16 = sizeof(ggml_fp16_t) * ne; -+ const size_t qx_sz = ne * ggml_type_size(quant)/ggml_blck_size(quant); -+ float * x = (float *) malloc(x_sz); -+ void * qx = malloc(qx_sz); -+ vk_buffer qx_buf = ggml_vk_create_buffer_check(ctx->device, qx_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); -+ vk_buffer x_buf = ggml_vk_create_buffer_check(ctx->device, x_sz_f16, vk::MemoryPropertyFlagBits::eDeviceLocal); -+ float * x_ref = (float *) malloc(x_sz); -+ ggml_fp16_t * x_chk = (ggml_fp16_t *) malloc(x_sz_f16); -+ -+ for (size_t i = 0; i < ne; i++) { -+ x[i] = rand() / (float)RAND_MAX; -+ } -+ -+ vk_pipeline p = ggml_vk_get_to_fp16(ctx, quant); -+ -+ ggml_vk_quantize_data(x, qx, ne, quant); -+ ggml_vk_dequantize_data(qx, x_ref, ne, quant); -+ -+ ggml_pipeline_request_descriptor_sets(ctx->device, p, 1); -+ -+ ggml_pipeline_allocate_descriptor_sets(ctx->device); -+ -+ ggml_vk_buffer_write(qx_buf, 0, qx, qx_sz); -+ -+ vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); -+ ggml_vk_ctx_begin(ctx->device, subctx); -+ const std::vector pc = { 1, (uint32_t)ne, (uint32_t)ne, (uint32_t)ne, (uint32_t)ne }; -+ ggml_vk_dispatch_pipeline(ctx, subctx, p, { vk_subbuffer{ qx_buf, 0, qx_sz }, vk_subbuffer{ x_buf, 0, x_sz_f16 } }, pc.size() * sizeof(int), pc.data(), { (uint32_t)ne, 1, 1}); -+ ggml_vk_ctx_end(subctx); -+ -+ auto begin = std::chrono::high_resolution_clock::now(); -+ -+ ggml_vk_submit(subctx, ctx->fence); -+ VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences"); -+ ctx->device->device.resetFences({ ctx->fence }); -+ -+ auto end = std::chrono::high_resolution_clock::now(); -+ -+ double ms_dequant = std::chrono::duration_cast(end-begin).count() / 1000.0; -+ ggml_vk_buffer_read(x_buf, 0, x_chk, x_sz_f16); -+ -+ int first_err = -1; -+ -+ double avg_err = 0.0; -+ for (size_t i = 0; i < ne; i++) { -+ double error = std::fabs(x_ref[i] - ggml_fp16_to_fp32(x_chk[i])); -+ avg_err += error; -+ -+ if (first_err < 0 && error > 0.05) { -+ first_err = i; -+ } -+ } -+ -+ avg_err /= ne; -+ -+ std::cerr << "TEST DEQUANT " << ggml_type_name(quant) << " time=" << ms_dequant << "ms avg_err=" << avg_err << std::endl; -+ -+ if (avg_err > 0.1) { -+ std::cerr << "first_error = " << first_err << std::endl; -+ std::cerr << "Actual result: " << std::endl << std::endl; -+ for (int i = std::max(0, first_err - 5); i < std::min((int)ne, first_err + 5); i++) { -+ std::cerr << ggml_fp16_to_fp32(x_chk[i]) << ", "; -+ } -+ std::cerr << std::endl << "Expected result: " << std::endl << std::endl; -+ for (int i = std::max(0, first_err - 5); i < std::min((int)ne, first_err + 5); i++) { -+ std::cerr << x_ref[i] << ", "; -+ } -+ std::cerr << std::endl; -+ } -+ -+ ggml_vk_destroy_buffer(x_buf); -+ ggml_vk_destroy_buffer(qx_buf); -+ -+ free(x); -+ free(qx); -+ free(x_ref); -+ free(x_chk); -+} -+ -+static void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, size_t split_k, size_t shader_size, ggml_type quant) { -+ VK_LOG_DEBUG("ggml_vk_test_dequant_matmul(" << m << ", " << n << ", " << k << ", " << batch << ", " << num_it << ", " << split_k << ", " << ggml_type_name(quant) << ")"); -+ const size_t x_ne = m * k * batch; -+ const size_t y_ne = k * n * batch; -+ const size_t d_ne = m * n * batch; -+ -+ vk_pipeline p; -+ std::string shname; -+ if (shader_size == 0) { -+ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_s : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_s; -+ shname = std::string(ggml_type_name(quant)) + "_ALIGNED_S"; -+ } else if (shader_size == 1) { -+ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_m : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_m; -+ shname = std::string(ggml_type_name(quant)) + "_ALIGNED_M"; -+ } else if (shader_size == 2) { -+ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->a_l : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->a_l; -+ shname = std::string(ggml_type_name(quant)) + "_ALIGNED_L"; -+ } else { -+ GGML_ASSERT(0); -+ } -+ -+ const size_t kpad = ggml_vk_align_size(k, p->align); -+ -+ if (k != kpad) { -+ if (shader_size == 0) { -+ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->s : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->s; -+ shname = std::string(ggml_type_name(quant)) + "_S"; -+ } else if (shader_size == 1) { -+ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->m : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->m; -+ shname = std::string(ggml_type_name(quant)) + "_M"; -+ } else if (shader_size == 2) { -+ p = ctx->device->fp16 ? ctx->device->pipeline_dequant_mul_mat_mat[quant].f16acc->l : ctx->device->pipeline_dequant_mul_mat_mat[quant].f32acc->l; -+ shname = std::string(ggml_type_name(quant)) + "_L"; -+ } else { -+ GGML_ASSERT(0); -+ } -+ } -+ -+ const size_t x_sz = sizeof(float) * x_ne; -+ const size_t y_sz = sizeof(float) * y_ne; -+ const size_t qx_sz = x_ne * ggml_type_size(quant)/ggml_blck_size(quant); -+ const size_t d_sz = sizeof(float) * d_ne; -+ float * x = (float *) malloc(x_sz); -+ float * y = (float *) malloc(y_sz); -+ void * qx = malloc(qx_sz); -+ vk_buffer qx_buf = ggml_vk_create_buffer_check(ctx->device, qx_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); -+ vk_buffer y_buf = ggml_vk_create_buffer_check(ctx->device, y_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); -+ vk_buffer d_buf = ggml_vk_create_buffer_check(ctx->device, d_sz, vk::MemoryPropertyFlagBits::eDeviceLocal); -+ float * d = (float *) malloc(d_sz); -+ float * d_chk = (float *) malloc(d_sz); -+ -+ for (size_t i = 0; i < x_ne; i++) { -+ x[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f; -+ } -+ -+ ggml_vk_quantize_data(x, qx, x_ne, quant); -+ -+ for (size_t i = 0; i < y_ne; i++) { -+ // y[i] = rand() / (float)RAND_MAX; -+ y[i] = (i % k == i / k) ? 1.0f : 0.0f; -+ } -+ -+ ggml_pipeline_request_descriptor_sets(ctx->device, p, num_it); -+ if (split_k > 1) { -+ ggml_pipeline_request_descriptor_sets(ctx->device, ctx->device->pipeline_matmul_split_k_reduce, num_it); -+ -+ if (ctx->prealloc_split_k == nullptr || ctx->prealloc_split_k->size < sizeof(float) * d_ne * split_k) { -+ // Resize buffer -+ if (ctx->prealloc_split_k != nullptr) { -+ ggml_vk_destroy_buffer(ctx->prealloc_split_k); -+ } -+ ctx->prealloc_split_k = ggml_vk_create_buffer_check(ctx->device, sizeof(float) * d_ne * split_k, vk::MemoryPropertyFlagBits::eDeviceLocal); -+ } -+ } -+ -+ ggml_pipeline_allocate_descriptor_sets(ctx->device); -+ -+ ggml_vk_buffer_write(qx_buf, 0, qx, qx_sz); -+ ggml_vk_buffer_write(y_buf, 0, y, y_sz); -+ -+ vk_context subctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); -+ ggml_vk_ctx_begin(ctx->device, subctx); -+ for (size_t i = 0; i < num_it; i++) { -+ ggml_vk_matmul( -+ ctx, subctx, p, ggml_vk_subbuffer(qx_buf), ggml_vk_subbuffer(y_buf), ggml_vk_subbuffer(d_buf), ggml_vk_subbuffer(ctx->prealloc_split_k), -+ m, n, k, -+ k, k, m, k*m, k*n, m*n, -+ split_k, batch, batch, batch, 1, 1 -+ ); -+ } -+ ggml_vk_ctx_end(subctx); -+ -+ auto begin = std::chrono::high_resolution_clock::now(); -+ -+ ggml_vk_submit(subctx, ctx->fence); -+ VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences"); -+ ctx->device->device.resetFences({ ctx->fence }); -+ -+ auto end = std::chrono::high_resolution_clock::now(); -+ -+ double time_ms = std::chrono::duration_cast(end-begin).count() / 1000.0; -+ ggml_vk_buffer_read(d_buf, 0, d, d_sz); -+ -+ ggml_init_params iparams = { -+ /*.mem_size =*/ 1024*1024*1024, -+ /*.mem_buffer =*/ NULL, -+ /*.no_alloc =*/ true, -+ }; -+ -+ ggml_context * ggml_ctx = ggml_init(iparams); -+ -+ ggml_tensor * src0_ggml = ggml_new_tensor_3d(ggml_ctx, quant, k, m, batch); -+ ggml_tensor * src1_ggml = ggml_new_tensor_3d(ggml_ctx, GGML_TYPE_F32, k, n, batch); -+ ggml_tensor * tensor_ggml = ggml_mul_mat(ggml_ctx, src0_ggml, src1_ggml); -+ -+ src0_ggml->data = qx; -+ src1_ggml->data = y; -+ tensor_ggml->data = d_chk; -+ -+ ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); -+ ggml_build_forward_expand(cgraph, tensor_ggml); -+ -+ ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 1); -+ -+ ggml_free(ggml_ctx); -+ -+ double avg_err = 0.0; -+ int first_err_n = -1; -+ int first_err_m = -1; -+ int first_err_b = -1; -+ -+ for (size_t i = 0; i < m*n*batch; i++) { -+ double err = std::fabs(d[i] - d_chk[i]); -+ avg_err += err; -+ -+ if ((err > 0.05f || std::isnan(err)) && first_err_n == -1) { -+ first_err_b = i / (m * n); -+ first_err_n = (i % (m * n)) / m; -+ first_err_m = (i % (m * n)) % m; -+ } -+ } -+ -+ avg_err /= m * n; -+ -+ double tflops = 2.0*m*n*k*batch*num_it / (time_ms / 1000.0) / (1000.0*1000.0*1000.0*1000.0); -+ -+ std::cerr << "TEST MMQ " << shname << " m=" << m << " n=" << n << " k=" << k << " batch=" << batch << " split_k=" << split_k << " matmul " << time_ms / num_it << "ms " << tflops << " TFLOPS avg_err=" << avg_err << std::endl; -+ -+ if (avg_err > 0.01 || std::isnan(avg_err)) { -+ std::cerr << "m = " << first_err_m << " n = " << first_err_n << " b = " << first_err_b << std::endl; -+ std::cerr << "Actual result: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(d, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ std::cerr << std::endl; -+ std::cerr << "Expected result: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(d_chk, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ -+ if (split_k > 1) { -+ float * split_k_buf = (float *) malloc(sizeof(float) * d_ne * split_k); -+ ggml_vk_buffer_read(ctx->prealloc_split_k, 0, split_k_buf, sizeof(float) * d_ne * split_k); -+ -+ std::cerr << "d_buf0: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(split_k_buf, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ -+ std::cerr << "d_buf1: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(split_k_buf + d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ -+ std::cerr << "d_buf2: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(split_k_buf + 2 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ -+ std::cerr << "d_buf3: " << std::endl << std::endl; -+ ggml_vk_print_matrix_area(split_k_buf + 3 * d_ne, GGML_TYPE_F32, m, n, first_err_m, first_err_n, first_err_b); -+ -+ free(split_k_buf); -+ } -+ } -+ -+ ggml_vk_destroy_buffer(qx_buf); -+ ggml_vk_destroy_buffer(y_buf); -+ ggml_vk_destroy_buffer(d_buf); -+ -+ free(x); -+ free(qx); -+ free(y); -+ free(d); -+ free(d_chk); -+} -+#endif -+ -+static void ggml_vk_preallocate_buffers(ggml_backend_vk_context * ctx) { -+#if defined(GGML_VULKAN_RUN_TESTS) -+ const std::vector vals { -+ 512, 512, 128, -+ 128, 512, 512, -+ 4096, 512, 4096, -+ 11008, 512, 4096, -+ 4096, 512, 11008, -+ 32000, 512, 4096, -+ 8, 8, 8, -+ 100, 46, 576, -+ 623, 111, 128, -+ 100, 46, 558, -+ 512, 1, 256, -+ 128, 110, 622, -+ 511, 511, 127, -+ 511, 511, 7, -+ 511, 511, 17, -+ 49, 49, 128, -+ 128, 49, 49, -+ 4096, 49, 4096, -+ }; -+ const size_t num_it = 100; -+ -+ for (size_t i = 0; i < vals.size(); i += 3) { -+ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0); -+ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1); -+ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2); -+ std::cerr << '\n'; -+ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0); -+ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1); -+ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2); -+ std::cerr << '\n'; -+ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0); -+ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1); -+ ggml_vk_test_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2); -+ std::cerr << '\n' << std::endl; -+ -+ if (vals[i + 2] % 32 == 0) { -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0, GGML_TYPE_Q4_0); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1, GGML_TYPE_Q4_0); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2, GGML_TYPE_Q4_0); -+ std::cerr << '\n'; -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0, GGML_TYPE_Q4_0); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1, GGML_TYPE_Q4_0); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2, GGML_TYPE_Q4_0); -+ std::cerr << '\n'; -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0, GGML_TYPE_Q4_0); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1, GGML_TYPE_Q4_0); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2, GGML_TYPE_Q4_0); -+ std::cerr << '\n' << std::endl; -+ } -+ -+ if (vals[i + 2] % 256 == 0) { -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 0, GGML_TYPE_Q4_K); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 1, GGML_TYPE_Q4_K); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 1, 2, GGML_TYPE_Q4_K); -+ std::cerr << '\n'; -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 0, GGML_TYPE_Q4_K); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 1, GGML_TYPE_Q4_K); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 2, 2, GGML_TYPE_Q4_K); -+ std::cerr << '\n'; -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 0, GGML_TYPE_Q4_K); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 1, GGML_TYPE_Q4_K); -+ ggml_vk_test_dequant_matmul(ctx, vals[i], vals[i + 1], vals[i + 2], 2, num_it, 4, 2, GGML_TYPE_Q4_K); -+ std::cerr << '\n' << std::endl; -+ } -+ } -+ -+ GGML_ABORT("fatal error"); -+#endif -+ -+ if (ctx->prealloc_x == nullptr || (ctx->prealloc_size_x > 0 && ctx->prealloc_x->size < ctx->prealloc_size_x)) { -+ VK_LOG_MEMORY("ggml_vk_preallocate_buffers(x_size: " << ctx->prealloc_size_x << ")"); -+ // Resize buffer -+ if (ctx->prealloc_x != nullptr) { -+ ggml_vk_destroy_buffer(ctx->prealloc_x); -+ } -+ ctx->prealloc_x = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_x); -+ } -+ if (ctx->prealloc_y == nullptr || (ctx->prealloc_size_y > 0 && ctx->prealloc_y->size < ctx->prealloc_size_y)) { -+ VK_LOG_MEMORY("ggml_vk_preallocate_buffers(y_size: " << ctx->prealloc_size_y << ")"); -+ // Resize buffer -+ if (ctx->prealloc_y != nullptr) { -+ ggml_vk_destroy_buffer(ctx->prealloc_y); -+ } -+ ctx->prealloc_y = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_y); -+ } -+ if (ctx->prealloc_split_k == nullptr || (ctx->prealloc_size_split_k > 0 && ctx->prealloc_split_k->size < ctx->prealloc_size_split_k)) { -+ VK_LOG_MEMORY("ggml_vk_preallocate_buffers(split_k_size: " << ctx->prealloc_size_split_k << ")"); -+ // Resize buffer -+ if (ctx->prealloc_split_k != nullptr) { -+ ggml_vk_destroy_buffer(ctx->prealloc_split_k); -+ } -+ ctx->prealloc_split_k = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_split_k); -+ } -+} -+ -+static bool ggml_vk_compute_forward(ggml_backend_vk_context* ctx, ggml_tensor* tensor, int tensor_idx, bool use_fence); -+ -+// Returns true if node has enqueued work into the queue, false otherwise -+// If submit is true the current all operations queued so far are being submitted to Vulkan to overlap cmdlist creation and GPU execution. -+static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_tensor * node, int node_idx, ggml_tensor *node_begin, int node_idx_begin, bool dryrun, bool last_node, bool submit){ -+ if (ggml_is_empty(node) || !node->buffer) { -+ return false; -+ } -+ -+ VK_LOG_DEBUG("ggml_vk_build_graph(" << node << ", " << ggml_op_name(node->op) << ")"); -+ ctx->semaphore_idx = 0; -+ -+ const ggml_tensor * src0 = node->src[0]; -+ const ggml_tensor * src1 = node->src[1]; -+ const ggml_tensor * src2 = node->src[2]; -+ const ggml_tensor * src3 = node->src[3]; -+ -+ switch (node->op) { -+ // Return on empty ops to avoid generating a compute_ctx and setting exit_tensor -+ case GGML_OP_RESHAPE: -+ case GGML_OP_VIEW: -+ case GGML_OP_PERMUTE: -+ case GGML_OP_TRANSPOSE: -+ case GGML_OP_NONE: -+ return false; -+ case GGML_OP_UNARY: -+ switch (ggml_get_unary_op(node)) { -+ case GGML_UNARY_OP_SILU: -+ case GGML_UNARY_OP_GELU: -+ case GGML_UNARY_OP_GELU_QUICK: -+ case GGML_UNARY_OP_RELU: -+ case GGML_UNARY_OP_TANH: -+ break; -+ default: -+ return false; -+ } -+ break; -+ case GGML_OP_REPEAT: -+ case GGML_OP_GET_ROWS: -+ case GGML_OP_ADD: -+ case GGML_OP_ACC: -+ case GGML_OP_MUL: -+ case GGML_OP_DIV: -+ case GGML_OP_CONCAT: -+ case GGML_OP_UPSCALE: -+ case GGML_OP_SCALE: -+ case GGML_OP_SQR: -+ case GGML_OP_SIN: -+ case GGML_OP_COS: -+ case GGML_OP_CLAMP: -+ case GGML_OP_PAD: -+ case GGML_OP_CPY: -+ case GGML_OP_CONT: -+ case GGML_OP_DUP: -+ case GGML_OP_NORM: -+ case GGML_OP_GROUP_NORM: -+ case GGML_OP_RMS_NORM: -+ case GGML_OP_DIAG_MASK_INF: -+ case GGML_OP_SOFT_MAX: -+ case GGML_OP_ROPE: -+ case GGML_OP_MUL_MAT: -+ case GGML_OP_MUL_MAT_ID: -+ case GGML_OP_ARGSORT: -+ case GGML_OP_SUM_ROWS: -+ case GGML_OP_IM2COL: -+ case GGML_OP_TIMESTEP_EMBEDDING: -+ case GGML_OP_POOL_2D: -+ case GGML_OP_RWKV_WKV6: -+ case GGML_OP_LEAKY_RELU: -+ case GGML_OP_FLASH_ATTN_EXT: -+ break; -+ default: -+ std::cerr << "ggml_vulkan: Error: Missing op: " << ggml_op_name(node->op) << std::endl; -+ GGML_ABORT("fatal error"); -+ return false; -+ } -+ -+ vk_context compute_ctx; -+ -+ if (!dryrun) { -+ if (ctx->compute_ctx.expired()) { -+ compute_ctx = ggml_vk_create_context(ctx, ctx->device->compute_queue); -+ ctx->compute_ctx = compute_ctx; -+ ggml_vk_ctx_begin(ctx->device, compute_ctx); -+ } else { -+ compute_ctx = ctx->compute_ctx.lock(); -+ } -+ } else { -+ switch (node->op) { -+ case GGML_OP_REPEAT: -+ case GGML_OP_ACC: -+ case GGML_OP_GET_ROWS: -+ case GGML_OP_ADD: -+ case GGML_OP_MUL: -+ case GGML_OP_DIV: -+ case GGML_OP_CONCAT: -+ case GGML_OP_UPSCALE: -+ case GGML_OP_SCALE: -+ case GGML_OP_SQR: -+ case GGML_OP_SIN: -+ case GGML_OP_COS: -+ case GGML_OP_CLAMP: -+ case GGML_OP_PAD: -+ case GGML_OP_CPY: -+ case GGML_OP_CONT: -+ case GGML_OP_DUP: -+ case GGML_OP_NORM: -+ case GGML_OP_GROUP_NORM: -+ case GGML_OP_RMS_NORM: -+ case GGML_OP_UNARY: -+ case GGML_OP_DIAG_MASK_INF: -+ case GGML_OP_SOFT_MAX: -+ case GGML_OP_ROPE: -+ case GGML_OP_ARGSORT: -+ case GGML_OP_SUM_ROWS: -+ case GGML_OP_IM2COL: -+ case GGML_OP_TIMESTEP_EMBEDDING: -+ case GGML_OP_POOL_2D: -+ case GGML_OP_LEAKY_RELU: -+ { -+ // These operations all go through ggml_vk_op_f32, so short-circuit and -+ // do the only thing needed for the dryrun. -+ vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, src2, node, node->op); -+ ggml_pipeline_request_descriptor_sets(ctx->device, pipeline, 1); -+ return false; -+ } -+ default: -+ break; -+ } -+ } -+ -+ switch (node->op) { -+ case GGML_OP_REPEAT: -+ ggml_vk_repeat(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_ACC: -+ ggml_vk_acc(ctx, compute_ctx, src0, src1, node, dryrun); -+ -+ break; -+ case GGML_OP_GET_ROWS: -+ ggml_vk_get_rows(ctx, compute_ctx, src0, src1, node, dryrun); -+ -+ break; -+ case GGML_OP_ADD: -+ ggml_vk_add(ctx, compute_ctx, src0, src1, node, dryrun); -+ -+ break; -+ case GGML_OP_MUL: -+ ggml_vk_mul(ctx, compute_ctx, src0, src1, node, dryrun); -+ -+ break; -+ case GGML_OP_DIV: -+ ggml_vk_div(ctx, compute_ctx, src0, src1, node, dryrun); -+ -+ break; -+ case GGML_OP_CONCAT: -+ ggml_vk_concat(ctx, compute_ctx, src0, src1, node, dryrun); -+ -+ break; -+ case GGML_OP_UPSCALE: -+ ggml_vk_upscale(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_SCALE: -+ ggml_vk_scale(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_SQR: -+ ggml_vk_sqr(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_SIN: -+ ggml_vk_sin(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_COS: -+ ggml_vk_cos(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_CLAMP: -+ ggml_vk_clamp(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_PAD: -+ ggml_vk_pad(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_CPY: -+ case GGML_OP_CONT: -+ case GGML_OP_DUP: -+ ggml_vk_cpy(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_NORM: -+ ggml_vk_norm(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_GROUP_NORM: -+ ggml_vk_group_norm(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_RMS_NORM: -+ ggml_vk_rms_norm(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_UNARY: -+ switch (ggml_get_unary_op(node)) { -+ case GGML_UNARY_OP_SILU: -+ case GGML_UNARY_OP_GELU: -+ case GGML_UNARY_OP_GELU_QUICK: -+ case GGML_UNARY_OP_RELU: -+ case GGML_UNARY_OP_TANH: -+ ggml_vk_unary(ctx, compute_ctx, src0, node, dryrun); -+ break; -+ default: -+ return false; -+ } -+ break; -+ case GGML_OP_DIAG_MASK_INF: -+ ggml_vk_diag_mask_inf(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_SOFT_MAX: -+ ggml_vk_soft_max(ctx, compute_ctx, src0, src1, node, dryrun); -+ -+ break; -+ case GGML_OP_ROPE: -+ ggml_vk_rope(ctx, compute_ctx, src0, src1, src2, node, dryrun); -+ -+ break; -+ case GGML_OP_ARGSORT: -+ ggml_vk_argsort(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_SUM_ROWS: -+ ggml_vk_sum_rows(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_IM2COL: -+ ggml_vk_im2col(ctx, compute_ctx, src0, src1, node, dryrun); -+ -+ break; -+ case GGML_OP_TIMESTEP_EMBEDDING: -+ ggml_vk_timestep_embedding(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_POOL_2D: -+ ggml_vk_pool_2d(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_LEAKY_RELU: -+ ggml_vk_leaky_relu(ctx, compute_ctx, src0, node, dryrun); -+ -+ break; -+ case GGML_OP_MUL_MAT: -+ ggml_vk_mul_mat(ctx, compute_ctx, src0, src1, node, dryrun); -+ -+ break; -+ case GGML_OP_MUL_MAT_ID: -+ ggml_vk_mul_mat_id(ctx, compute_ctx, src0, src1, src2, node, dryrun); -+ -+ break; -+ -+ case GGML_OP_FLASH_ATTN_EXT: -+ ggml_vk_flash_attn(ctx, compute_ctx, src0, src1, src2, src3, node, dryrun); -+ -+ break; -+ -+ case GGML_OP_RWKV_WKV6: -+ ggml_vk_rwkv_wkv6(ctx, compute_ctx, node, dryrun); -+ -+ break; -+ default: -+ return false; -+ } -+ -+ if (dryrun) { -+ return false; -+ } -+ -+ ctx->tensor_ctxs[node_idx] = compute_ctx; -+ -+#if defined(GGML_VULKAN_CHECK_RESULTS) || defined(GGML_VULKAN_PERF) -+ // Force context reset on each node so that each tensor ends up in its own context -+ // and can be run and compared to its CPU equivalent separately -+ last_node = true; -+#endif -+ -+ if (submit || last_node) { -+ ggml_vk_ctx_end(compute_ctx); -+ -+ // TODO probably it'd be better to pass a exit_node flag to ggml_vk_compute_forward -+ if (last_node) { -+ compute_ctx->exit_tensor_idx = node_idx_begin; -+ } -+ else { -+ compute_ctx->exit_tensor_idx = -1; -+ } -+ -+ ctx->compute_ctx.reset(); -+ -+ bool ok = ggml_vk_compute_forward(ctx, node_begin, node_idx_begin, false); -+ if (!ok) { -+ if (node->op == GGML_OP_UNARY) { -+ std::cerr << __func__ << ": error: op not supported UNARY " << node->name << " (" << ggml_unary_op_name(static_cast(node->op_params[0])) << ")" << std::endl; -+ } -+ else { -+ std::cerr << __func__ << ": error: op not supported " << node->name << " (" << ggml_op_name(node->op) << ")" << std::endl; -+ } -+ } -+ -+ } -+ return true; -+} -+ -+static bool ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_tensor * tensor, int tensor_idx, bool use_fence = true){ -+ ggml_backend_buffer * buf = nullptr; -+ -+ switch (tensor->op) { -+ case GGML_OP_ADD: -+ case GGML_OP_ACC: -+ case GGML_OP_GET_ROWS: -+ case GGML_OP_MUL: -+ case GGML_OP_DIV: -+ case GGML_OP_CONCAT: -+ case GGML_OP_UPSCALE: -+ case GGML_OP_SCALE: -+ case GGML_OP_SQR: -+ case GGML_OP_SIN: -+ case GGML_OP_COS: -+ case GGML_OP_CLAMP: -+ case GGML_OP_PAD: -+ case GGML_OP_CPY: -+ case GGML_OP_CONT: -+ case GGML_OP_DUP: -+ case GGML_OP_NORM: -+ case GGML_OP_GROUP_NORM: -+ case GGML_OP_RMS_NORM: -+ case GGML_OP_DIAG_MASK_INF: -+ case GGML_OP_SOFT_MAX: -+ case GGML_OP_ROPE: -+ case GGML_OP_RESHAPE: -+ case GGML_OP_VIEW: -+ case GGML_OP_PERMUTE: -+ case GGML_OP_TRANSPOSE: -+ case GGML_OP_NONE: -+ case GGML_OP_ARGSORT: -+ case GGML_OP_SUM_ROWS: -+ case GGML_OP_IM2COL: -+ case GGML_OP_TIMESTEP_EMBEDDING: -+ case GGML_OP_POOL_2D: -+ case GGML_OP_RWKV_WKV6: -+ case GGML_OP_LEAKY_RELU: -+ case GGML_OP_REPEAT: -+ buf = tensor->buffer; -+ -+ break; -+ case GGML_OP_UNARY: -+ switch (ggml_get_unary_op(tensor)) { -+ case GGML_UNARY_OP_SILU: -+ case GGML_UNARY_OP_GELU: -+ case GGML_UNARY_OP_GELU_QUICK: -+ case GGML_UNARY_OP_RELU: -+ case GGML_UNARY_OP_TANH: -+ buf = tensor->buffer; -+ break; -+ default: -+ return false; -+ } -+ break; -+ case GGML_OP_MUL_MAT: -+ case GGML_OP_MUL_MAT_ID: -+ case GGML_OP_FLASH_ATTN_EXT: -+ buf = tensor->buffer; -+ -+ break; -+ default: -+ return false; -+ } -+ -+ if (buf == nullptr) { -+ return false; -+ } -+ -+ VK_LOG_DEBUG("ggml_vk_compute_forward(" << tensor << ", name=" << tensor->name << ", op=" << ggml_op_name(tensor->op) << ", type=" << tensor->type << ", ne0=" << tensor->ne[0] << ", ne1=" << tensor->ne[1] << ", ne2=" << tensor->ne[2] << ", ne3=" << tensor->ne[3] << ", nb0=" << tensor->nb[0] << ", nb1=" << tensor->nb[1] << ", nb2=" << tensor->nb[2] << ", nb3=" << tensor->nb[3] << ", view_src=" << tensor->view_src << ", view_offs=" << tensor->view_offs << ")"); -+ -+ vk_context subctx = ctx->tensor_ctxs[tensor_idx].lock(); -+ -+ // always wait for the GPU work to be done for the last submit -+ if (tensor_idx == subctx->exit_tensor_idx) { -+ use_fence = true; -+ } -+ -+ // Only run if ctx hasn't been submitted yet -+ if (!subctx->seqs.empty()) { -+#ifdef GGML_VULKAN_CHECK_RESULTS -+ ggml_vk_check_results_0(tensor); -+ use_fence = true; -+#endif -+ -+ // Do staging buffer copies -+ for (auto& cpy : subctx->in_memcpys) { -+ memcpy(cpy.dst, cpy.src, cpy.n); -+ } -+ -+ ggml_vk_submit(subctx, use_fence ? ctx->fence : vk::Fence{}); -+ -+ if (use_fence) { -+ VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_compute_forward waitForFences"); -+ -+ ctx->device->device.resetFences({ ctx->fence }); -+ } -+#ifdef GGML_VULKAN_CHECK_RESULTS -+ ggml_vk_check_results_1(tensor); -+#endif -+ } -+ -+ if (tensor_idx == subctx->exit_tensor_idx) { -+ // Do staging buffer copies -+ for (auto& cpy : subctx->out_memcpys) { -+ memcpy(cpy.dst, cpy.src, cpy.n); -+ } -+ subctx->in_memcpys.clear(); -+ subctx->out_memcpys.clear(); -+ } -+ -+ return true; -+} -+ -+// Clean up after graph processing is done -+static void ggml_vk_graph_cleanup(ggml_backend_vk_context * ctx) { -+ VK_LOG_DEBUG("ggml_vk_graph_cleanup()"); -+ for (auto& buffer : ctx->gc.temp_buffers) { -+ ggml_vk_pool_free(ctx, buffer); -+ } -+ ctx->gc.temp_buffers.clear(); -+ -+ for (auto& dsr : ctx->device->pipeline_descriptor_set_requirements) { -+ vk_pipeline_ref plr = ctx->device->pipelines[dsr.first]; -+ -+ if (plr.expired()) { -+ continue; -+ } -+ -+ vk_pipeline pl = plr.lock(); -+ ggml_pipeline_cleanup(pl); -+ } -+ -+ ggml_vk_queue_cleanup(ctx->device, ctx->device->compute_queue); -+ ggml_vk_queue_cleanup(ctx->device, ctx->device->transfer_queue); -+ -+ for (size_t i = 0; i < ctx->gc.semaphores.size(); i++) { -+ ctx->device->device.destroySemaphore({ ctx->gc.semaphores[i].s }); -+ } -+ ctx->gc.semaphores.clear(); -+ -+ for (size_t i = 0; i < ctx->gc.tl_semaphores.size(); i++) { -+ ctx->device->device.destroySemaphore({ ctx->gc.tl_semaphores[i].s }); -+ } -+ ctx->gc.tl_semaphores.clear(); -+ ctx->semaphore_idx = 0; -+ -+ ctx->event_idx = 0; -+ -+ for (auto& event : ctx->gc.events) { -+ ctx->device->device.resetEvent(event); -+ } -+ -+ ctx->tensor_ctxs.clear(); -+ ctx->gc.contexts.clear(); -+ ctx->device->pipeline_descriptor_set_requirements.clear(); -+} -+ -+// Clean up on backend free -+static void ggml_vk_cleanup(ggml_backend_vk_context * ctx) { -+ VK_LOG_DEBUG("ggml_vk_cleanup(" << ctx->name << ")"); -+ ggml_vk_graph_cleanup(ctx); -+ -+ ggml_vk_destroy_buffer(ctx->prealloc_x); -+ ggml_vk_destroy_buffer(ctx->prealloc_y); -+ ggml_vk_destroy_buffer(ctx->prealloc_split_k); -+ -+ for (auto& buffer : ctx->buffer_pool) { -+ ggml_vk_destroy_buffer(buffer); -+ } -+ -+ ctx->prealloc_size_x = 0; -+ ctx->prealloc_size_y = 0; -+ ctx->prealloc_size_split_k = 0; -+ -+ for (auto& event : ctx->gc.events) { -+ ctx->device->device.destroyEvent(event); -+ } -+ ctx->gc.events.clear(); -+ -+ ctx->device->device.destroyFence(ctx->fence); -+} -+ -+static int ggml_vk_get_device_count() { -+ ggml_vk_instance_init(); -+ -+ return vk_instance.device_indices.size(); -+} -+ -+static void ggml_vk_get_device_description(int device, char * description, size_t description_size) { -+ ggml_vk_instance_init(); -+ -+ std::vector devices = vk_instance.instance.enumeratePhysicalDevices(); -+ -+ vk::PhysicalDeviceProperties props; -+ devices[device].getProperties(&props); -+ -+ snprintf(description, description_size, "%s", props.deviceName.data()); -+} -+ -+// backend interface -+ -+#define UNUSED GGML_UNUSED -+ -+// device backend -+ -+static bool ggml_backend_buffer_is_vk(ggml_backend_buffer_t buffer) { -+ return buffer->buft->iface.get_name == ggml_backend_vk_buffer_type_name; -+} -+ -+static void ggml_backend_vk_buffer_free_buffer(ggml_backend_buffer_t buffer) { -+ VK_LOG_MEMORY("ggml_backend_vk_buffer_free_buffer()"); -+ ggml_backend_vk_buffer_context * ctx = (ggml_backend_vk_buffer_context *)buffer->context; -+ ggml_vk_destroy_buffer(ctx->dev_buffer); -+ delete ctx; -+} -+ -+static void * ggml_backend_vk_buffer_get_base(ggml_backend_buffer_t buffer) { -+ return vk_ptr_base; -+ -+ UNUSED(buffer); -+} -+ -+static void ggml_backend_vk_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { -+ VK_LOG_DEBUG("ggml_backend_vk_buffer_init_tensor(" << buffer << " (" << buffer->context << "), " << tensor << ")"); -+ if (tensor->view_src != nullptr) { -+ GGML_ASSERT(tensor->view_src->buffer->buft == buffer->buft); -+ } -+} -+ -+static void ggml_backend_vk_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { -+ VK_LOG_DEBUG("ggml_backend_vk_buffer_set_tensor(" << buffer << ", " << tensor << ", " << data << ", " << offset << ", " << size << ")"); -+ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)buffer->context; -+ vk_buffer buf = buf_ctx->dev_buffer; -+ -+ ggml_vk_buffer_write(buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); -+} -+ -+static void ggml_backend_vk_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { -+ VK_LOG_DEBUG("ggml_backend_vk_buffer_get_tensor(" << buffer << ", " << tensor << ", " << data << ", " << offset << ", " << size << ")"); -+ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)buffer->context; -+ -+ vk_buffer buf = buf_ctx->dev_buffer; -+ -+ ggml_vk_buffer_read(buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); -+} -+ -+static bool ggml_backend_vk_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) { -+ if (ggml_backend_buffer_is_vk(src->buffer)) { -+ ggml_backend_vk_buffer_context * src_buf_ctx = (ggml_backend_vk_buffer_context *)src->buffer->context; -+ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; -+ -+ vk_buffer src_buf = src_buf_ctx->dev_buffer; -+ vk_buffer dst_buf = dst_buf_ctx->dev_buffer; -+ -+ ggml_vk_buffer_copy(dst_buf, vk_tensor_offset(dst) + dst->view_offs, src_buf, vk_tensor_offset(src) + src->view_offs, ggml_nbytes(src)); -+ -+ return true; -+ } -+ return false; -+ -+ UNUSED(buffer); -+} -+ -+static void ggml_backend_vk_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { -+ ggml_backend_vk_buffer_context * ctx = (ggml_backend_vk_buffer_context *)buffer->context; -+ -+ ggml_vk_buffer_memset(ctx->dev_buffer, 0, value, buffer->size); -+} -+ -+static ggml_backend_buffer_i ggml_backend_vk_buffer_interface = { -+ /* .free_buffer = */ ggml_backend_vk_buffer_free_buffer, -+ /* .get_base = */ ggml_backend_vk_buffer_get_base, -+ /* .init_tensor = */ ggml_backend_vk_buffer_init_tensor, -+ /* .memset_tensor = */ NULL, -+ /* .set_tensor = */ ggml_backend_vk_buffer_set_tensor, -+ /* .get_tensor = */ ggml_backend_vk_buffer_get_tensor, -+ /* .cpy_tensor = */ ggml_backend_vk_buffer_cpy_tensor, -+ /* .clear = */ ggml_backend_vk_buffer_clear, -+ /* .reset = */ NULL, -+}; -+ -+// vk buffer type -+static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft) { -+ ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *)buft->context; -+ -+ return ctx->name.c_str(); -+} -+ -+static ggml_backend_buffer_t ggml_backend_vk_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { -+ VK_LOG_MEMORY("ggml_backend_vk_buffer_type_alloc_buffer(" << size << ")"); -+ ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; -+ -+ vk_buffer dev_buffer = nullptr; -+ try { -+ dev_buffer = ggml_vk_create_buffer_device(ctx->device, size); -+ } catch (const vk::SystemError& e) { -+ return nullptr; -+ } -+ -+ ggml_backend_vk_buffer_context * bufctx = new ggml_backend_vk_buffer_context(ctx->device, std::move(dev_buffer), ctx->name); -+ -+ return ggml_backend_buffer_init(buft, ggml_backend_vk_buffer_interface, bufctx, size); -+} -+ -+static size_t ggml_backend_vk_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { -+ ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; -+ return ctx->device->properties.limits.minStorageBufferOffsetAlignment; -+} -+ -+static size_t ggml_backend_vk_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { -+ ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; -+ return ctx->device->max_memory_allocation_size; -+} -+ -+static size_t ggml_backend_vk_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { -+ return ggml_nbytes(tensor); -+ -+ UNUSED(buft); -+} -+ -+ggml_backend_buffer_type_t ggml_backend_vk_buffer_type(size_t dev_num) { -+ ggml_vk_instance_init(); -+ -+ VK_LOG_DEBUG("ggml_backend_vk_buffer_type(" << dev_num << ")"); -+ -+ vk_device dev = ggml_vk_get_device(dev_num); -+ -+ return &dev->buffer_type; -+} -+ -+// host buffer type -+ -+static const char * ggml_backend_vk_host_buffer_type_name(ggml_backend_buffer_type_t buft) { -+ return GGML_VK_NAME "_Host"; -+ -+ UNUSED(buft); -+} -+ -+static const char * ggml_backend_vk_host_buffer_name(ggml_backend_buffer_t buffer) { -+ return GGML_VK_NAME "_Host"; -+ -+ UNUSED(buffer); -+} -+ -+static void ggml_backend_vk_host_buffer_free_buffer(ggml_backend_buffer_t buffer) { -+ VK_LOG_MEMORY("ggml_backend_vk_host_buffer_free_buffer()"); -+ ggml_vk_host_free(vk_instance.devices[0], buffer->context); -+} -+ -+static ggml_backend_buffer_t ggml_backend_vk_host_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { -+ VK_LOG_MEMORY("ggml_backend_vk_host_buffer_type_alloc_buffer(" << size << ")"); -+ -+ size += 32; // Behave like the CPU buffer type -+ void * ptr = nullptr; -+ try { -+ ptr = ggml_vk_host_malloc(vk_instance.devices[0], size); -+ } catch (vk::SystemError& e) { -+ std::cerr << "ggml_vulkan: Failed to allocate pinned memory." << std::endl; -+ std::cerr << "ggml_vulkan: " << e.what() << std::endl; -+ // fallback to cpu buffer -+ return ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); -+ } -+ -+ ggml_backend_buffer_t buffer = ggml_backend_cpu_buffer_from_ptr(ptr, size); -+ buffer->buft = buft; -+ buffer->iface.free_buffer = ggml_backend_vk_host_buffer_free_buffer; -+ -+ return buffer; -+ -+ UNUSED(buft); -+} -+ -+static size_t ggml_backend_vk_host_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { -+ return vk_instance.devices[0]->properties.limits.minMemoryMapAlignment; -+ -+ UNUSED(buft); -+} -+ -+// Should be changed to return device-specific host buffer type -+// but that probably requires changes in llama.cpp -+ggml_backend_buffer_type_t ggml_backend_vk_host_buffer_type() { -+ static struct ggml_backend_buffer_type ggml_backend_vk_buffer_type_host = { -+ /* .iface = */ { -+ /* .get_name = */ ggml_backend_vk_host_buffer_type_name, -+ /* .alloc_buffer = */ ggml_backend_vk_host_buffer_type_alloc_buffer, -+ /* .get_alignment = */ ggml_backend_vk_host_buffer_type_get_alignment, -+ /* .get_max_size = */ NULL, // defaults to SIZE_MAX -+ /* .get_alloc_size = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size, -+ /* .is_host = */ ggml_backend_cpu_buffer_type()->iface.is_host, -+ }, -+ /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), 0), -+ /* .context = */ nullptr, -+ }; -+ -+ // Make sure device 0 is initialized -+ ggml_vk_instance_init(); -+ ggml_vk_get_device(0); -+ -+ return &ggml_backend_vk_buffer_type_host; -+} -+ -+ -+// backend -+ -+static const char * ggml_backend_vk_name(ggml_backend_t backend) { -+ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; -+ -+ return ctx->name.c_str(); -+} -+ -+static void ggml_backend_vk_free(ggml_backend_t backend) { -+ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; -+ VK_LOG_DEBUG("ggml_backend_vk_free(" << ctx->name << ")"); -+ -+ ggml_vk_cleanup(ctx); -+ -+ delete ctx; -+ delete backend; -+} -+ -+static ggml_backend_buffer_type_t ggml_backend_vk_get_default_buffer_type(ggml_backend_t backend) { -+ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; -+ -+ return &ctx->device->buffer_type; -+} -+ -+static void ggml_backend_vk_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { -+ VK_LOG_DEBUG("ggml_backend_vk_set_tensor_async(" << size << ")"); -+ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; -+ GGML_ASSERT((tensor->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || tensor->buffer->buft == ggml_backend_vk_host_buffer_type()) && "unsupported buffer type"); -+ -+ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; -+ -+ vk_context transfer_ctx; -+ -+ if (ctx->transfer_ctx.expired()) { -+ // Initialize new transfer context -+ transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); -+ ctx->transfer_ctx = transfer_ctx; -+ ggml_vk_ctx_begin(ctx->device, transfer_ctx); -+ } else { -+ transfer_ctx = ctx->transfer_ctx.lock(); -+ } -+ -+ vk_buffer buf = buf_ctx->dev_buffer; -+ -+ ggml_vk_buffer_write_async(transfer_ctx, buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); -+} -+ -+static void ggml_backend_vk_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { -+ VK_LOG_DEBUG("ggml_backend_vk_get_tensor_async(" << size << ")"); -+ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; -+ GGML_ASSERT((tensor->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || tensor->buffer->buft == ggml_backend_vk_host_buffer_type()) && "unsupported buffer type"); -+ -+ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; -+ -+ vk_context transfer_ctx; -+ -+ if (ctx->transfer_ctx.expired()) { -+ // Initialize new transfer context -+ transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); -+ ctx->transfer_ctx = transfer_ctx; -+ ggml_vk_ctx_begin(ctx->device, transfer_ctx); -+ } else { -+ transfer_ctx = ctx->transfer_ctx.lock(); -+ } -+ -+ vk_buffer buf = buf_ctx->dev_buffer; -+ -+ ggml_vk_buffer_read_async(transfer_ctx, buf, vk_tensor_offset(tensor) + tensor->view_offs + offset, data, size); -+} -+ -+static bool ggml_backend_vk_cpy_tensor_async(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst) { -+ VK_LOG_DEBUG("ggml_backend_vk_cpy_tensor_async()"); -+ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; -+ if ((dst->buffer->buft == ggml_backend_vk_get_default_buffer_type(backend) || dst->buffer->buft == ggml_backend_vk_host_buffer_type()) && ggml_backend_buffer_is_vk(src->buffer)) { -+ ggml_backend_vk_buffer_context * src_buf_ctx = (ggml_backend_vk_buffer_context *)src->buffer->context; -+ ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; -+ -+ vk_context transfer_ctx; -+ -+ if (ctx->transfer_ctx.expired()) { -+ // Initialize new transfer context -+ transfer_ctx = ggml_vk_create_context(ctx, ctx->device->transfer_queue); -+ ctx->transfer_ctx = transfer_ctx; -+ ggml_vk_ctx_begin(ctx->device, transfer_ctx); -+ } else { -+ transfer_ctx = ctx->transfer_ctx.lock(); -+ } -+ -+ vk_buffer src_buf = src_buf_ctx->dev_buffer; -+ vk_buffer dst_buf = dst_buf_ctx->dev_buffer; -+ -+ ggml_vk_buffer_copy_async(transfer_ctx, dst_buf, vk_tensor_offset(dst) + dst->view_offs, src_buf, vk_tensor_offset(src) + src->view_offs, ggml_nbytes(src)); -+ return true; -+ } -+ -+ return false; -+} -+ -+static void ggml_backend_vk_synchronize(ggml_backend_t backend) { -+ VK_LOG_DEBUG("ggml_backend_vk_synchronize()"); -+ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; -+ if(ctx->transfer_ctx.expired()) { -+ return; -+ } -+ -+ vk_context transfer_ctx = ctx->transfer_ctx.lock(); -+ -+ ggml_vk_ctx_end(transfer_ctx); -+ -+ for (auto& cpy : transfer_ctx->in_memcpys) { -+ memcpy(cpy.dst, cpy.src, cpy.n); -+ } -+ -+ ggml_vk_submit(transfer_ctx, ctx->fence); -+ VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_backend_vk_synchronize waitForFences"); -+ ctx->device->device.resetFences({ ctx->fence }); -+ -+ for (auto& cpy : transfer_ctx->out_memcpys) { -+ memcpy(cpy.dst, cpy.src, cpy.n); -+ } -+ -+ ctx->transfer_ctx.reset(); -+} -+ -+static bool ggml_vk_is_empty(ggml_tensor * node) { -+ return ggml_is_empty(node) || node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE; -+} -+ -+static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { -+ VK_LOG_DEBUG("ggml_backend_vk_graph_compute(" << cgraph->n_nodes << " nodes)"); -+ ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; -+ -+ for (int i = 0; i < cgraph->n_nodes; i++) { -+ ggml_vk_build_graph(ctx, cgraph->nodes[i], i, nullptr, 0, true, false, false); -+ } -+ ggml_vk_preallocate_buffers(ctx); -+ ggml_pipeline_allocate_descriptor_sets(ctx->device); -+ -+ int last_node = cgraph->n_nodes - 1; -+ -+ // If the last op in the cgraph isn't backend GPU, the command buffer doesn't get closed properly -+ while (last_node > 0 && ggml_vk_is_empty(cgraph->nodes[last_node])) { -+ last_node -= 1; -+ } -+ -+ // Reserve tensor context space for all nodes -+ ctx->tensor_ctxs.resize(cgraph->n_nodes); -+ -+ bool first_node_in_batch = true; // true if next node will be first node in a batch -+ int submit_node_idx = 0; // index to first node in a batch -+ -+ // Submit work every nodes_per_submit nodes to overlap CPU cmdbuffer generation with GPU execution. -+ // Start with a smaller count to get work submitted right away, and increase it after each submit. -+ int nodes_per_submit = 20; -+ int submitted_nodes = 0; -+ int submit_count = 0; -+ for (int i = 0; i < cgraph->n_nodes; i++) { -+ if (first_node_in_batch) { -+ submit_node_idx = i; -+ } -+ -+ bool submit = (submitted_nodes >= nodes_per_submit) || (i == last_node); -+ -+ bool enqueued = ggml_vk_build_graph(ctx, cgraph->nodes[i], i, cgraph->nodes[submit_node_idx], submit_node_idx, false, i == last_node, submit); -+ -+ if (enqueued) { -+ ++submitted_nodes; -+ -+#ifndef GGML_VULKAN_CHECK_RESULTS -+ if (first_node_in_batch) { -+ first_node_in_batch = false; -+ } -+#endif -+ } -+ -+ if (submit) { -+ first_node_in_batch = true; -+ submitted_nodes = 0; -+ switch (submit_count) { -+ case 0: -+ nodes_per_submit = 50; -+ break; -+ default: -+ nodes_per_submit = 100; -+ break; -+ } -+ submit_count++; -+ } -+ } -+ -+#ifdef GGML_VULKAN_PERF -+ ctx->device->perf_logger->print_timings(); -+#endif -+ -+ ggml_vk_graph_cleanup(ctx); -+ -+ return GGML_STATUS_SUCCESS; -+ -+ UNUSED(backend); -+} -+ -+// TODO: enable async and synchronize -+static ggml_backend_i ggml_backend_vk_interface = { -+ /* .get_name = */ ggml_backend_vk_name, -+ /* .free = */ ggml_backend_vk_free, -+ /* .set_tensor_async = */ NULL, // ggml_backend_vk_set_tensor_async, -+ /* .get_tensor_async = */ NULL, // ggml_backend_vk_get_tensor_async, -+ /* .cpy_tensor_async = */ NULL, // ggml_backend_vk_cpy_tensor_async, -+ /* .synchronize = */ NULL, // ggml_backend_vk_synchronize, -+ /* .graph_plan_create = */ NULL, -+ /* .graph_plan_free = */ NULL, -+ /* .graph_plan_update = */ NULL, -+ /* .graph_plan_compute = */ NULL, -+ /* .graph_compute = */ ggml_backend_vk_graph_compute, -+ /* .event_record = */ NULL, -+ /* .event_wait = */ NULL, -+}; -+ -+static ggml_guid_t ggml_backend_vk_guid() { -+ static ggml_guid guid = { 0xb8, 0xf7, 0x4f, 0x86, 0x40, 0x3c, 0xe1, 0x02, 0x91, 0xc8, 0xdd, 0xe9, 0x02, 0x3f, 0xc0, 0x2b }; -+ return &guid; -+} -+ -+ggml_backend_t ggml_backend_vk_init(size_t dev_num) { -+ VK_LOG_DEBUG("ggml_backend_vk_init(" << dev_num << ")"); -+ -+ ggml_backend_vk_context * ctx = new ggml_backend_vk_context; -+ ggml_vk_init(ctx, dev_num); -+ -+ ggml_backend_t vk_backend = new ggml_backend { -+ /* .guid = */ ggml_backend_vk_guid(), -+ /* .interface = */ ggml_backend_vk_interface, -+ /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), dev_num), -+ /* .context = */ ctx, -+ }; -+ -+ return vk_backend; -+} -+ -+bool ggml_backend_is_vk(ggml_backend_t backend) { -+ return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_vk_guid()); -+} -+ -+int ggml_backend_vk_get_device_count() { -+ return ggml_vk_get_device_count(); -+} -+ -+void ggml_backend_vk_get_device_description(int device, char * description, size_t description_size) { -+ GGML_ASSERT(device < (int) vk_instance.device_indices.size()); -+ int dev_idx = vk_instance.device_indices[device]; -+ ggml_vk_get_device_description(dev_idx, description, description_size); -+} -+ -+void ggml_backend_vk_get_device_memory(int device, size_t * free, size_t * total) { -+ GGML_ASSERT(device < (int) vk_instance.device_indices.size()); -+ -+ vk::PhysicalDevice vkdev = vk_instance.instance.enumeratePhysicalDevices()[vk_instance.device_indices[device]]; -+ -+ vk::PhysicalDeviceMemoryProperties memprops = vkdev.getMemoryProperties(); -+ -+ for (const vk::MemoryHeap& heap : memprops.memoryHeaps) { -+ if (heap.flags & vk::MemoryHeapFlagBits::eDeviceLocal) { -+ *total = heap.size; -+ *free = heap.size; -+ break; -+ } -+ } -+} -+ -+////////////////////////// -+ -+struct ggml_backend_vk_device_context { -+ size_t device; -+ std::string name; -+ std::string description; -+}; -+ -+static const char * ggml_backend_vk_device_get_name(ggml_backend_dev_t dev) { -+ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; -+ return ctx->name.c_str(); -+} -+ -+static const char * ggml_backend_vk_device_get_description(ggml_backend_dev_t dev) { -+ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; -+ return ctx->description.c_str(); -+} -+ -+static void ggml_backend_vk_device_get_memory(ggml_backend_dev_t device, size_t * free, size_t * total) { -+ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)device->context; -+ ggml_backend_vk_get_device_memory(ctx->device, free, total); -+} -+ -+static ggml_backend_buffer_type_t ggml_backend_vk_device_get_buffer_type(ggml_backend_dev_t dev) { -+ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; -+ return ggml_backend_vk_buffer_type(ctx->device); -+} -+ -+static ggml_backend_buffer_type_t ggml_backend_vk_device_get_host_buffer_type(ggml_backend_dev_t dev) { -+ UNUSED(dev); -+ return ggml_backend_vk_host_buffer_type(); -+} -+ -+static enum ggml_backend_dev_type ggml_backend_vk_device_get_type(ggml_backend_dev_t dev) { -+ UNUSED(dev); -+ return GGML_BACKEND_DEVICE_TYPE_GPU; -+} -+ -+static void ggml_backend_vk_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props) { -+ props->name = ggml_backend_vk_device_get_name(dev); -+ props->description = ggml_backend_vk_device_get_description(dev); -+ props->type = ggml_backend_vk_device_get_type(dev); -+ ggml_backend_vk_device_get_memory(dev, &props->memory_free, &props->memory_total); -+ props->caps = { -+ /* .async = */ false, -+ /* .host_buffer = */ true, -+ /* .buffer_from_host_ptr = */ false, -+ /* .events = */ false, -+ }; -+} -+ -+static ggml_backend_t ggml_backend_vk_device_init(ggml_backend_dev_t dev, const char * params) { -+ UNUSED(params); -+ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; -+ return ggml_backend_vk_init(ctx->device); -+} -+ -+static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { -+ switch (op->op) { -+ case GGML_OP_UNARY: -+ switch (ggml_get_unary_op(op)) { -+ case GGML_UNARY_OP_GELU: -+ case GGML_UNARY_OP_GELU_QUICK: -+ case GGML_UNARY_OP_SILU: -+ case GGML_UNARY_OP_RELU: -+ case GGML_UNARY_OP_TANH: -+ return ggml_is_contiguous(op->src[0]); -+ default: -+ return false; -+ } -+ break; -+ case GGML_OP_MUL_MAT: -+ case GGML_OP_MUL_MAT_ID: -+ { -+ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; -+ const vk_device& device = ggml_vk_get_device(ctx->device); -+ if (op->op == GGML_OP_MUL_MAT_ID && !device->mul_mat_id_s && !device->mul_mat_id_m && !device->mul_mat_id_l) { -+ // If there's not enough shared memory for row_ids and the result tile, fallback to CPU -+ return false; -+ } -+ switch (op->src[0]->type) { -+ case GGML_TYPE_F32: -+ case GGML_TYPE_F16: -+ case GGML_TYPE_Q4_0: -+ case GGML_TYPE_Q4_1: -+ case GGML_TYPE_Q5_0: -+ case GGML_TYPE_Q5_1: -+ case GGML_TYPE_Q8_0: -+ case GGML_TYPE_Q2_K: -+ case GGML_TYPE_Q3_K: -+ case GGML_TYPE_Q4_K: -+ case GGML_TYPE_Q5_K: -+ case GGML_TYPE_Q6_K: -+ case GGML_TYPE_IQ4_NL: -+ break; -+ default: -+ return false; -+ } -+ struct ggml_tensor * a; -+ struct ggml_tensor * b; -+ if (op->op == GGML_OP_MUL_MAT) { -+ a = op->src[0]; -+ b = op->src[1]; -+ } else { -+ a = op->src[2]; -+ b = op->src[1]; -+ } -+ if (a->ne[3] != b->ne[3]) { -+ return false; -+ } -+ if (!(ggml_vk_dim01_contiguous(op->src[0]) || op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) || -+ !(ggml_vk_dim01_contiguous(op->src[1]) || op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16)) { -+ return false; -+ } -+ -+ return true; -+ } break; -+ case GGML_OP_FLASH_ATTN_EXT: -+ { -+ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; -+ if (!ggml_vk_get_device(ctx->device)->coopmat2) { -+ return false; -+ } -+ switch (op->src[0]->ne[0]) { -+ case 64: -+ case 80: -+ case 96: -+ case 112: -+ case 128: -+ case 256: -+ break; -+ default: -+ return false; -+ } -+ if (op->src[0]->type != GGML_TYPE_F32) { -+ return false; -+ } -+ if (op->type != GGML_TYPE_F32) { -+ return false; -+ } -+ if (op->src[3] && op->src[3]->type != GGML_TYPE_F16) { -+ return false; -+ } -+ // It's straightforward to support different K/V dequant, but would -+ // significantly increase the number of pipelines -+ if (op->src[1]->type != op->src[2]->type) { -+ return false; -+ } -+ switch (op->src[1]->type) { -+ case GGML_TYPE_F16: -+ case GGML_TYPE_Q4_0: -+ case GGML_TYPE_Q4_1: -+ case GGML_TYPE_Q5_0: -+ case GGML_TYPE_Q5_1: -+ case GGML_TYPE_Q8_0: -+ // K dequants currently disabled because D dimension is rounded up to 256 and runs inefficiently -+ //case GGML_TYPE_Q2_K: -+ //case GGML_TYPE_Q3_K: -+ //case GGML_TYPE_Q4_K: -+ //case GGML_TYPE_Q5_K: -+ //case GGML_TYPE_Q6_K: -+ case GGML_TYPE_IQ4_NL: -+ break; -+ default: -+ return false; -+ } -+ return true; -+ } -+ case GGML_OP_GET_ROWS: -+ { -+ switch (op->src[0]->type) { -+ case GGML_TYPE_F32: -+ case GGML_TYPE_F16: -+ case GGML_TYPE_Q4_0: -+ case GGML_TYPE_Q4_1: -+ case GGML_TYPE_Q5_0: -+ case GGML_TYPE_Q5_1: -+ case GGML_TYPE_Q8_0: -+ case GGML_TYPE_IQ4_NL: -+ return true; -+ default: -+ return false; -+ } -+ } break; -+ case GGML_OP_CONT: -+ case GGML_OP_CPY: -+ case GGML_OP_DUP: -+ { -+ ggml_type src0_type = op->src[0]->type; -+ ggml_type src1_type = op->src[1] != nullptr ? op->src[1]->type : src0_type; -+ if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { -+ return true; -+ } -+ if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F16) { -+ return true; -+ } -+ if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { -+ return true; -+ } -+ return false; -+ } break; -+ case GGML_OP_REPEAT: -+ return ggml_type_size(op->type) == sizeof(float) && ggml_type_size(op->src[0]->type) == sizeof(float); -+ case GGML_OP_ROPE: -+ { -+ const int mode = ((const int32_t *) op->op_params)[2]; -+ if (mode & GGML_ROPE_TYPE_MROPE) { -+ return false; -+ } -+ if (mode & GGML_ROPE_TYPE_VISION) { -+ return false; -+ } -+ return ggml_is_contiguous(op->src[0]); -+ } -+ case GGML_OP_NONE: -+ case GGML_OP_RESHAPE: -+ case GGML_OP_VIEW: -+ case GGML_OP_PERMUTE: -+ case GGML_OP_TRANSPOSE: -+ case GGML_OP_NORM: -+ case GGML_OP_GROUP_NORM: -+ case GGML_OP_RMS_NORM: -+ case GGML_OP_ADD: -+ case GGML_OP_ACC: -+ case GGML_OP_MUL: -+ case GGML_OP_DIV: -+ case GGML_OP_CONCAT: -+ case GGML_OP_UPSCALE: -+ case GGML_OP_SCALE: -+ case GGML_OP_SQR: -+ case GGML_OP_SIN: -+ case GGML_OP_COS: -+ case GGML_OP_CLAMP: -+ case GGML_OP_PAD: -+ case GGML_OP_DIAG_MASK_INF: -+ case GGML_OP_SOFT_MAX: -+ case GGML_OP_ARGSORT: -+ case GGML_OP_SUM_ROWS: -+ case GGML_OP_IM2COL: -+ case GGML_OP_TIMESTEP_EMBEDDING: -+ case GGML_OP_POOL_2D: -+ case GGML_OP_RWKV_WKV6: -+ case GGML_OP_LEAKY_RELU: -+ return true; -+ default: -+ return false; -+ } -+ -+ UNUSED(dev); -+} -+ -+static bool ggml_backend_vk_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { -+ if (buft->iface.get_name != ggml_backend_vk_buffer_type_name) { -+ return false; -+ } -+ -+ ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; -+ ggml_backend_vk_buffer_type_context * buft_ctx = (ggml_backend_vk_buffer_type_context *)buft->context; -+ -+ return buft_ctx->device->idx == ctx->device; -+} -+ -+static bool ggml_backend_vk_device_offload_op(ggml_backend_dev_t dev, const ggml_tensor * op) { -+ const int min_batch_size = 32; -+ -+ return (op->ne[1] >= min_batch_size && op->op != GGML_OP_GET_ROWS) || -+ (op->ne[2] >= min_batch_size && op->op == GGML_OP_MUL_MAT_ID); -+ -+ UNUSED(dev); -+} -+ -+static const struct ggml_backend_device_i ggml_backend_vk_device_i = { -+ /* .get_name = */ ggml_backend_vk_device_get_name, -+ /* .get_description = */ ggml_backend_vk_device_get_description, -+ /* .get_memory = */ ggml_backend_vk_device_get_memory, -+ /* .get_type = */ ggml_backend_vk_device_get_type, -+ /* .get_props = */ ggml_backend_vk_device_get_props, -+ /* .init_backend = */ ggml_backend_vk_device_init, -+ /* .get_buffer_type = */ ggml_backend_vk_device_get_buffer_type, -+ /* .get_host_buffer_type = */ ggml_backend_vk_device_get_host_buffer_type, -+ /* .buffer_from_host_ptr = */ NULL, -+ /* .supports_op = */ ggml_backend_vk_device_supports_op, -+ /* .supports_buft = */ ggml_backend_vk_device_supports_buft, -+ /* .offload_op = */ ggml_backend_vk_device_offload_op, -+ /* .event_new = */ NULL, -+ /* .event_free = */ NULL, -+ /* .event_synchronize = */ NULL, -+}; -+ -+static const char * ggml_backend_vk_reg_get_name(ggml_backend_reg_t reg) { -+ UNUSED(reg); -+ return GGML_VK_NAME; -+} -+ -+static size_t ggml_backend_vk_reg_get_device_count(ggml_backend_reg_t reg) { -+ UNUSED(reg); -+ return ggml_backend_vk_get_device_count(); -+} -+ -+static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg, size_t device) { -+ static std::vector devices; -+ -+ static bool initialized = false; -+ -+ { -+ static std::mutex mutex; -+ std::lock_guard lock(mutex); -+ if (!initialized) { -+ for (int i = 0; i < ggml_backend_vk_get_device_count(); i++) { -+ ggml_backend_vk_device_context * ctx = new ggml_backend_vk_device_context; -+ char desc[256]; -+ ggml_backend_vk_get_device_description(i, desc, sizeof(desc)); -+ ctx->device = i; -+ ctx->name = GGML_VK_NAME + std::to_string(i); -+ ctx->description = desc; -+ devices.push_back(new ggml_backend_device { -+ /* .iface = */ ggml_backend_vk_device_i, -+ /* .reg = */ reg, -+ /* .context = */ ctx, -+ }); -+ } -+ initialized = true; -+ } -+ } -+ -+ GGML_ASSERT(device < devices.size()); -+ return devices[device]; -+} -+ -+static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = { -+ /* .get_name = */ ggml_backend_vk_reg_get_name, -+ /* .get_device_count = */ ggml_backend_vk_reg_get_device_count, -+ /* .get_device = */ ggml_backend_vk_reg_get_device, -+ /* .get_proc_address = */ NULL, -+}; -+ -+ggml_backend_reg_t ggml_backend_vk_reg() { -+ static ggml_backend_reg reg = { -+ /* .api_version = */ GGML_BACKEND_API_VERSION, -+ /* .iface = */ ggml_backend_vk_reg_i, -+ /* .context = */ nullptr, -+ }; -+ -+ return ® -+} -+ -+// Extension availability -+static bool ggml_vk_instance_validation_ext_available(const std::vector& instance_extensions) { -+#ifdef GGML_VULKAN_VALIDATE -+ bool portability_enumeration_ext = false; -+ // Check for portability enumeration extension for MoltenVK support -+ for (const auto& properties : instance_extensions) { -+ if (strcmp("VK_KHR_portability_enumeration", properties.extensionName) == 0) { -+ return true; -+ } -+ } -+ if (!portability_enumeration_ext) { -+ std::cerr << "ggml_vulkan: WARNING: Instance extension VK_KHR_portability_enumeration not found." << std::endl; -+ } -+#endif -+ return false; -+ -+ UNUSED(instance_extensions); -+} -+static bool ggml_vk_instance_portability_enumeration_ext_available(const std::vector& instance_extensions) { -+#ifdef __APPLE__ -+ bool portability_enumeration_ext = false; -+ // Check for portability enumeration extension for MoltenVK support -+ for (const auto& properties : instance_extensions) { -+ if (strcmp("VK_KHR_portability_enumeration", properties.extensionName) == 0) { -+ return true; -+ } -+ } -+ if (!portability_enumeration_ext) { -+ std::cerr << "ggml_vulkan: WARNING: Instance extension VK_KHR_portability_enumeration not found." << std::endl; -+ } -+#endif -+ return false; -+ -+ UNUSED(instance_extensions); -+} -+ -+static bool ggml_vk_khr_cooperative_matrix_support(const vk::PhysicalDeviceProperties& props, const vk::PhysicalDeviceDriverProperties& driver_props) { -+ switch (props.vendorID) { -+ case VK_VENDOR_ID_INTEL: -+ // Intel drivers don't support coopmat properly yet -+ return false; -+ case VK_VENDOR_ID_AMD: -+ if (driver_props.driverID == vk::DriverId::eAmdProprietary || driver_props.driverID == vk::DriverId::eAmdOpenSource) { -+ // Workaround for AMD proprietary driver reporting support on all GPUs -+ const std::string name = props.deviceName; -+ return name.rfind("AMD Radeon RX 7", 0) == 0 || name.rfind("AMD Radeon(TM) RX 7", 0) == 0 || // RDNA 3 consumer GPUs -+ name.rfind("AMD Radeon PRO W7", 0) == 0 || name.rfind("AMD Radeon(TM) PRO W7", 0) == 0 || // RDNA 3 workstation GPUs -+ name.rfind("AMD Radeon 7", 0) == 0 || name.rfind("AMD Radeon(TM) 7", 0) == 0; // RDNA 3 APUs -+ } -+ return true; -+ default: -+ return true; -+ } -+} -+ -+// checks -+ -+#ifdef GGML_VULKAN_CHECK_RESULTS -+static void ggml_vk_print_graph_origin(const ggml_tensor * tensor, std::vector& done, int level = 0) { -+ if (std::find(done.begin(), done.end(), tensor) != done.end() || level > 10) { -+ return; -+ } -+ for (int j = 0; j < level; j++) { -+ std::cerr << " "; -+ } -+ std::cerr << ggml_op_name(tensor->op) << " gpu=" << (tensor->extra != nullptr) << std::endl; -+ -+ done.push_back(tensor); -+ -+ for (int i = 0; i < GGML_MAX_SRC; i++) { -+ if (tensor->src[i] != nullptr) { -+ ggml_vk_print_graph_origin(tensor->src[i], done, level + 1); -+ } -+ } -+} -+ -+static void ggml_vk_print_tensor_area(const ggml_tensor * tensor, const void * data, int i0, int i1, int i2, int i3) { -+ if (tensor->type != GGML_TYPE_F32 && tensor->type != GGML_TYPE_F16 && tensor->type != GGML_TYPE_I32) { -+ return; -+ } -+ i0 = std::max(i0, 5); -+ i1 = std::max(i1, 5); -+ i2 = std::max(i2, 0); -+ i3 = std::max(i3, 0); -+ fprintf(stderr, " "); -+ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { -+ fprintf(stderr, "%7d ", idx1); -+ } -+ fprintf(stderr, "\n"); -+ for (int idx0 = i0 - 5; idx0 < i0 + 5; idx0++) { -+ fprintf(stderr, "%7d: ", idx0); -+ for (int idx1 = i1 - 5; idx1 < i1 + 5; idx1++) { -+ if (idx0 >= 0 && idx0 < tensor->ne[0] && idx1 >= 0 && idx1 < tensor->ne[1] && i2 >= 0 && i2 < tensor->ne[2] && i3 >= 0 && i3 < tensor->ne[3]) { -+ float val; -+ if (tensor->type == GGML_TYPE_F32) { -+ val = *(const float *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); -+ } else if (tensor->type == GGML_TYPE_F16) { -+ val = ggml_fp16_to_fp32(*(const ggml_fp16_t *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0])); -+ } else if (tensor->type == GGML_TYPE_I32) { -+ val = *(const int32_t *) ((const char *) data + i3*tensor->nb[3] + i2*tensor->nb[2] + idx1*tensor->nb[1] + idx0*tensor->nb[0]); -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ fprintf(stderr, "% 7.2f ", val); -+ } else { -+ fprintf(stderr, " "); -+ } -+ } -+ fprintf(stderr, "\n"); -+ } -+} -+ -+static void ggml_vk_print_tensor(const ggml_tensor * tensor, const char * name) { -+ void * tensor_data = tensor->data; -+ -+ const bool is_gpu = tensor->buffer != nullptr && ggml_backend_buffer_is_vk(tensor->buffer); -+ -+ if (is_gpu) { -+ const size_t tensor_size = ggml_nbytes(tensor); -+ tensor_data = malloc(tensor_size); -+ -+ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; -+ -+ vk_buffer buffer_gpu = buf_ctx->dev_buffer; -+ ggml_vk_buffer_read(buffer_gpu, vk_tensor_offset(tensor) + tensor->view_offs, tensor_data, tensor_size); -+ } -+ -+ std::cerr << "TENSOR CHECK " << name << " (" << tensor->name << "): " << ggml_op_name(tensor->op) << std::endl; -+ std::cerr << "tensor=" << tensor << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << std::endl; -+ if (tensor->src[0] != nullptr) { -+ std::cerr << "tensor->src[0]=" << tensor->src[0] << " name=" << tensor->src[0]->name << " op=" << ggml_op_name(tensor->src[0]->op) << " type=" << ggml_type_name(tensor->src[0]->type) << " ne0=" << tensor->src[0]->ne[0] << " nb0=" << tensor->src[0]->nb[0] << " ne1=" << tensor->src[0]->ne[1] << " nb1=" << tensor->src[0]->nb[1] << " ne2=" << tensor->src[0]->ne[2] << " nb2=" << tensor->src[0]->nb[2] << " ne3=" << tensor->src[0]->ne[3] << " nb3=" << tensor->src[0]->nb[3] << std::endl; -+ } -+ if (tensor->src[1] != nullptr) { -+ std::cerr << "tensor->src[1]=" << tensor->src[1] << " name=" << tensor->src[1]->name << " op=" << ggml_op_name(tensor->src[1]->op) << " type=" << ggml_type_name(tensor->src[1]->type) << " ne0=" << tensor->src[1]->ne[0] << " nb0=" << tensor->src[1]->nb[0] << " ne1=" << tensor->src[1]->ne[1] << " nb1=" << tensor->src[1]->nb[1] << " ne2=" << tensor->src[1]->ne[2] << " nb2=" << tensor->src[1]->nb[2] << " ne3=" << tensor->src[1]->ne[3] << " nb3=" << tensor->src[1]->nb[3] << std::endl; -+ } -+ std::cerr << std::endl << "Result:" << std::endl; -+ ggml_vk_print_tensor_area(tensor, tensor_data, 5, 5, 0, 0); -+ std::cerr << std::endl; -+ std::vector done; -+ ggml_vk_print_graph_origin(tensor, done); -+ -+ if (is_gpu) { -+ free(tensor_data); -+ } -+} -+ -+void * comp_result; -+size_t comp_size; -+size_t comp_nb[GGML_MAX_DIMS]; -+size_t check_counter = 0; -+static void ggml_vk_check_results_0(ggml_tensor * tensor) { -+ if (tensor->op == GGML_OP_TRANSPOSE) { -+ return; -+ } -+ -+ check_counter++; -+ if (!(vk_output_tensor > 0 && vk_output_tensor == check_counter) && check_counter <= vk_skip_checks) { -+ return; -+ } -+ -+ VK_LOG_DEBUG("ggml_vk_check_results_0(" << tensor->name << ")"); -+ -+ ggml_tensor * src0 = tensor->src[0]; -+ ggml_tensor * src1 = tensor->src[1]; -+ ggml_tensor * src2 = tensor->src[2]; -+ ggml_tensor * src3 = tensor->src[3]; -+ -+ struct ggml_init_params iparams = { -+ /*.mem_size =*/ 2ul*1024ul*1024ul*1024ul, -+ /*.mem_buffer =*/ NULL, -+ /*.no_alloc =*/ false, -+ }; -+ -+ struct ggml_context * ggml_ctx = ggml_init(iparams); -+ -+ struct ggml_tensor * src0_clone = nullptr; -+ struct ggml_tensor * src1_clone = nullptr; -+ struct ggml_tensor * src2_clone = nullptr; -+ struct ggml_tensor * src3_clone = nullptr; -+ struct ggml_tensor * tensor_clone = nullptr; -+ -+ size_t src0_size; -+ size_t src1_size; -+ size_t src2_size; -+ size_t src3_size; -+ -+ void * src0_buffer = nullptr; -+ void * src1_buffer = nullptr; -+ void * src2_buffer = nullptr; -+ void * src3_buffer = nullptr; -+ -+ if (src0 != nullptr) { -+ src0_clone = ggml_dup_tensor(ggml_ctx, src0); -+ -+ src0_size = ggml_nbytes(src0); -+ -+ src0_buffer = malloc(src0_size); -+ src0_clone->data = src0_buffer; -+ if (ggml_backend_buffer_is_host(src0->buffer)) { -+ memcpy(src0_clone->data, src0->data, src0_size); -+ memcpy(src0_clone->nb, src0->nb, sizeof(size_t) * GGML_MAX_DIMS); -+ } else if (ggml_backend_buffer_is_vk(src0->buffer)) { -+ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; -+ vk_buffer& buffer_gpu = buf_ctx->dev_buffer; -+ uint64_t offset = vk_tensor_offset(src0) + src0->view_offs; -+ if (!ggml_is_contiguous(src0) && ggml_vk_dim01_contiguous(src0)) { -+ for (int i3 = 0; i3 < src0->ne[3]; i3++) { -+ for (int i2 = 0; i2 < src0->ne[2]; i2++) { -+ const int idx = i3*src0->ne[2] + i2; -+ ggml_vk_buffer_read(buffer_gpu, offset + idx * src0->nb[2], ((char *)src0_clone->data + idx * src0_clone->nb[2]), src0->ne[1] * src0->nb[1]); -+ } -+ } -+ -+ src0_clone->nb[0] = src0->nb[0]; -+ src0_clone->nb[1] = src0->nb[1]; -+ for (int i = 2; i < GGML_MAX_DIMS; i++) { -+ src0_clone->nb[i] = src0_clone->nb[i - 1]*src0_clone->ne[i - 1]; -+ } -+ } else { -+ if (offset + src0_size >= buffer_gpu->size) { -+ src0_size = buffer_gpu->size - offset; -+ } -+ ggml_vk_buffer_read(buffer_gpu, offset, src0_clone->data, src0_size); -+ memcpy(src0_clone->nb, src0->nb, sizeof(size_t) * GGML_MAX_DIMS); -+ } -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ -+ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { -+ ggml_vk_print_tensor(src0, "src0"); -+ } -+ } -+ if (src1 != nullptr) { -+ src1_clone = ggml_dup_tensor(ggml_ctx, src1); -+ -+ src1_size = ggml_nbytes(src1); -+ -+ src1_buffer = malloc(src1_size); -+ src1_clone->data = src1_buffer; -+ if (ggml_backend_buffer_is_host(src1->buffer)) { -+ memcpy(src1_clone->data, src1->data, src1_size); -+ memcpy(src1_clone->nb, src1->nb, sizeof(size_t) * GGML_MAX_DIMS); -+ } else if (ggml_backend_buffer_is_vk(src1->buffer)) { -+ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; -+ vk_buffer& buffer_gpu = buf_ctx->dev_buffer; -+ uint64_t offset = vk_tensor_offset(src1) + src1->view_offs; -+ if (!ggml_is_contiguous(src1) && ggml_vk_dim01_contiguous(src1)) { -+ for (int i3 = 0; i3 < src1->ne[3]; i3++) { -+ for (int i2 = 0; i2 < src1->ne[2]; i2++) { -+ const int idx = i3*src1->ne[2] + i2; -+ ggml_vk_buffer_read(buffer_gpu, offset + idx * src1->nb[2], ((char *)src1_clone->data + idx * src1_clone->nb[2]), src1->ne[1] * src1->nb[1]); -+ } -+ } -+ -+ src1_clone->nb[0] = src1->nb[0]; -+ src1_clone->nb[1] = src1->nb[1]; -+ for (int i = 2; i < GGML_MAX_DIMS; i++) { -+ src1_clone->nb[i] = src1_clone->nb[i - 1]*src1_clone->ne[i - 1]; -+ } -+ } else { -+ if (offset + src1_size >= buffer_gpu->size) { -+ src1_size = buffer_gpu->size - offset; -+ } -+ ggml_vk_buffer_read(buffer_gpu, offset, src1_clone->data, src1_size); -+ memcpy(src1_clone->nb, src1->nb, sizeof(size_t) * GGML_MAX_DIMS); -+ } -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ -+ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { -+ ggml_vk_print_tensor(src1, "src1"); -+ } -+ } -+ if (src2 != nullptr) { -+ src2_clone = ggml_dup_tensor(ggml_ctx, src2); -+ -+ src2_size = ggml_nbytes(src2); -+ -+ src2_buffer = malloc(src2_size); -+ src2_clone->data = src2_buffer; -+ if (ggml_backend_buffer_is_host(src2->buffer)) { -+ memcpy(src2_clone->data, src2->data, src2_size); -+ memcpy(src2_clone->nb, src2->nb, sizeof(size_t) * GGML_MAX_DIMS); -+ } else if (ggml_backend_buffer_is_vk(src2->buffer)) { -+ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src2->buffer->context; -+ vk_buffer& buffer_gpu = buf_ctx->dev_buffer; -+ uint64_t offset = vk_tensor_offset(src2) + src2->view_offs; -+ if (!ggml_is_contiguous(src2) && ggml_vk_dim01_contiguous(src2)) { -+ for (int i3 = 0; i3 < src2->ne[3]; i3++) { -+ for (int i2 = 0; i2 < src2->ne[2]; i2++) { -+ const int idx = i3*src2->ne[2] + i2; -+ ggml_vk_buffer_read(buffer_gpu, offset + idx * src2->nb[2], ((char *)src2_clone->data + idx * src2_clone->nb[2]), src2->ne[1] * src2->nb[1]); -+ } -+ } -+ -+ src2_clone->nb[0] = src2->nb[0]; -+ src2_clone->nb[1] = src2->nb[1]; -+ for (int i = 2; i < GGML_MAX_DIMS; i++) { -+ src2_clone->nb[i] = src2_clone->nb[i - 1]*src2_clone->ne[i - 1]; -+ } -+ } else { -+ if (offset + src2_size >= buffer_gpu->size) { -+ src2_size = buffer_gpu->size - offset; -+ } -+ ggml_vk_buffer_read(buffer_gpu, offset, src2_clone->data, src2_size); -+ memcpy(src2_clone->nb, src2->nb, sizeof(size_t) * GGML_MAX_DIMS); -+ } -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ -+ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { -+ ggml_vk_print_tensor(src2, "src2"); -+ } -+ } -+ if (src3 != nullptr) { -+ src3_clone = ggml_dup_tensor(ggml_ctx, src3); -+ -+ src3_size = ggml_nbytes(src3); -+ -+ src3_buffer = malloc(src3_size); -+ src3_clone->data = src3_buffer; -+ if (ggml_backend_buffer_is_host(src3->buffer)) { -+ memcpy(src3_clone->data, src3->data, src3_size); -+ memcpy(src3_clone->nb, src3->nb, sizeof(size_t) * GGML_MAX_DIMS); -+ } else if (ggml_backend_buffer_is_vk(src3->buffer)) { -+ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src3->buffer->context; -+ vk_buffer& buffer_gpu = buf_ctx->dev_buffer; -+ uint64_t offset = vk_tensor_offset(src3) + src3->view_offs; -+ if (!ggml_is_contiguous(src3) && ggml_vk_dim01_contiguous(src3)) { -+ for (int i3 = 0; i3 < src3->ne[3]; i3++) { -+ for (int i2 = 0; i2 < src3->ne[2]; i2++) { -+ const int idx = i3*src3->ne[2] + i2; -+ ggml_vk_buffer_read(buffer_gpu, offset + idx * src3->nb[2], ((char *)src3_clone->data + idx * src3_clone->nb[2]), src3->ne[1] * src3->nb[1]); -+ } -+ } -+ -+ src3_clone->nb[0] = src3->nb[0]; -+ src3_clone->nb[1] = src3->nb[1]; -+ for (int i = 2; i < GGML_MAX_DIMS; i++) { -+ src3_clone->nb[i] = src3_clone->nb[i - 1]*src3_clone->ne[i - 1]; -+ } -+ } else { -+ if (offset + src3_size >= buffer_gpu->size) { -+ src3_size = buffer_gpu->size - offset; -+ } -+ ggml_vk_buffer_read(buffer_gpu, offset, src3_clone->data, src3_size); -+ memcpy(src3_clone->nb, src3->nb, sizeof(size_t) * GGML_MAX_DIMS); -+ } -+ } else { -+ GGML_ABORT("fatal error"); -+ } -+ -+ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { -+ ggml_vk_print_tensor(src3, "src3"); -+ } -+ } -+ -+ if (tensor->op == GGML_OP_FLASH_ATTN_EXT) { -+ const float *params = (const float *)tensor->op_params; -+ tensor_clone = ggml_flash_attn_ext(ggml_ctx, src0_clone, src1_clone, src2_clone, src3_clone, params[0], params[1], params[2]); -+ } else if (tensor->op == GGML_OP_MUL_MAT) { -+ tensor_clone = ggml_mul_mat(ggml_ctx, src0_clone, src1_clone); -+ } else if (tensor->op == GGML_OP_MUL_MAT_ID) { -+ tensor_clone = ggml_mul_mat_id(ggml_ctx, src0_clone, src1_clone, src2_clone); -+ } else if (tensor->op == GGML_OP_MUL) { -+ tensor_clone = ggml_mul(ggml_ctx, src0_clone, src1_clone); -+ } else if (tensor->op == GGML_OP_DIV) { -+ tensor_clone = ggml_div(ggml_ctx, src0_clone, src1_clone); -+ } else if (tensor->op == GGML_OP_CONCAT) { -+ tensor_clone = ggml_concat(ggml_ctx, src0_clone, src1_clone, *(int *)tensor->op_params); -+ } else if (tensor->op == GGML_OP_UPSCALE) { -+ tensor_clone = ggml_upscale_ext(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); -+ } else if (tensor->op == GGML_OP_SCALE) { -+ tensor_clone = ggml_scale(ggml_ctx, src0_clone, ((float *)tensor->op_params)[0]); -+ } else if (tensor->op == GGML_OP_SQR) { -+ tensor_clone = ggml_sqr(ggml_ctx, src0_clone); -+ } else if (tensor->op == GGML_OP_SIN) { -+ tensor_clone = ggml_sin(ggml_ctx, src0_clone); -+ } else if (tensor->op == GGML_OP_COS) { -+ tensor_clone = ggml_cos(ggml_ctx, src0_clone); -+ } else if (tensor->op == GGML_OP_CLAMP) { -+ tensor_clone = ggml_clamp(ggml_ctx, src0_clone, ((float *)tensor->op_params)[0], ((float *)tensor->op_params)[1]); -+ } else if (tensor->op == GGML_OP_PAD) { -+ tensor_clone = ggml_pad(ggml_ctx, src0_clone, tensor->ne[0] - src0_clone->ne[0], tensor->ne[1] - src0_clone->ne[1], tensor->ne[2] - src0_clone->ne[2], tensor->ne[3] - src0_clone->ne[3]); -+ } else if (tensor->op == GGML_OP_REPEAT) { -+ tensor_clone = ggml_repeat(ggml_ctx, src0_clone, tensor); -+ } else if (tensor->op == GGML_OP_ADD) { -+ tensor_clone = ggml_add(ggml_ctx, src0_clone, src1_clone); -+ } else if (tensor->op == GGML_OP_ACC) { -+ tensor_clone = ggml_acc(ggml_ctx, src0_clone, src1_clone, tensor->op_params[0], tensor->op_params[1], tensor->op_params[2], tensor->op_params[3]); -+ } else if (tensor->op == GGML_OP_NORM) { -+ tensor_clone = ggml_norm(ggml_ctx, src0_clone, *(float *)tensor->op_params); -+ } else if (tensor->op == GGML_OP_GROUP_NORM) { -+ tensor_clone = ggml_group_norm(ggml_ctx, src0_clone, *(int *)tensor->op_params, ((float *)tensor->op_params)[1]); -+ } else if (tensor->op == GGML_OP_RMS_NORM) { -+ tensor_clone = ggml_rms_norm(ggml_ctx, src0_clone, *(float *)tensor->op_params); -+ } else if (tensor->op == GGML_OP_SOFT_MAX) { -+ if (src1 != nullptr) { -+ tensor_clone = ggml_soft_max_ext(ggml_ctx, src0_clone, src1_clone, ((float *)tensor->op_params)[0], ((float *)tensor->op_params)[1]); -+ } else { -+ tensor_clone = ggml_soft_max(ggml_ctx, src0_clone); -+ } -+ } else if (tensor->op == GGML_OP_DIAG_MASK_INF) { -+ tensor_clone = ggml_diag_mask_inf(ggml_ctx, src0_clone, *(int *)tensor->op_params); -+ } else if (tensor->op == GGML_OP_ROPE) { -+ const int n_dims = ((int32_t *) tensor->op_params)[1]; -+ const int mode = ((int32_t *) tensor->op_params)[2]; -+ //const int n_ctx_ggml = ((int32_t *) tensor->op_params)[3]; -+ const int n_ctx_orig_ggml = ((int32_t *) tensor->op_params)[4]; -+ const float freq_base = ((float *) tensor->op_params)[5]; -+ const float freq_scale = ((float *) tensor->op_params)[6]; -+ const float ext_factor = ((float *) tensor->op_params)[7]; -+ const float attn_factor = ((float *) tensor->op_params)[8]; -+ const float beta_fast = ((float *) tensor->op_params)[9]; -+ const float beta_slow = ((float *) tensor->op_params)[10]; -+ tensor_clone = ggml_rope_ext(ggml_ctx, src0_clone, src1_clone, src2_clone, n_dims, mode, n_ctx_orig_ggml, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); -+ } else if (tensor->op == GGML_OP_UNARY) { -+ switch (ggml_get_unary_op(tensor)) { -+ case GGML_UNARY_OP_SILU: -+ tensor_clone = ggml_silu(ggml_ctx, src0_clone); -+ break; -+ case GGML_UNARY_OP_GELU: -+ tensor_clone = ggml_gelu(ggml_ctx, src0_clone); -+ break; -+ case GGML_UNARY_OP_GELU_QUICK: -+ tensor_clone = ggml_gelu_quick(ggml_ctx, src0_clone); -+ break; -+ case GGML_UNARY_OP_RELU: -+ tensor_clone = ggml_relu(ggml_ctx, src0_clone); -+ break; -+ case GGML_UNARY_OP_TANH: -+ tensor_clone = ggml_tanh(ggml_ctx, src0_clone); -+ break; -+ default: -+ std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl; -+ GGML_ABORT("fatal error"); -+ } -+ } else if (tensor->op == GGML_OP_CPY || tensor->op == GGML_OP_DUP) { -+ if (src1 == nullptr) { -+ tensor_clone = ggml_dup(ggml_ctx, src0_clone); -+ tensor_clone->type = tensor->type; -+ } else { -+ tensor_clone = ggml_cpy(ggml_ctx, src0_clone, src1_clone); -+ } -+ } else if (tensor->op == GGML_OP_CONT) { -+ tensor_clone = ggml_cont_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); -+ } else if (tensor->op == GGML_OP_RESHAPE) { -+ tensor_clone = ggml_reshape_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); -+ } else if (tensor->op == GGML_OP_VIEW) { -+ tensor_clone = ggml_view_4d(ggml_ctx, src0_clone, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3], tensor->nb[1], tensor->nb[2], tensor->nb[3], ((int32_t *) tensor->op_params)[0]); -+ } else if (tensor->op == GGML_OP_PERMUTE) { -+ int32_t * params = (int32_t *)tensor->op_params; -+ tensor_clone = ggml_permute(ggml_ctx, src0_clone, params[0], params[1], params[2], params[3]); -+ } else if (tensor->op == GGML_OP_TRANSPOSE) { -+ tensor_clone = ggml_transpose(ggml_ctx, src0_clone); -+ } else if (tensor->op == GGML_OP_GET_ROWS) { -+ tensor_clone = ggml_get_rows(ggml_ctx, src0_clone, src1_clone); -+ } else if (tensor->op == GGML_OP_ARGSORT) { -+ tensor_clone = ggml_argsort(ggml_ctx, src0_clone, (ggml_sort_order) *(int *)tensor->op_params); -+ } else if (tensor->op == GGML_OP_SUM_ROWS) { -+ tensor_clone = ggml_sum_rows(ggml_ctx, src0_clone); -+ } else if (tensor->op == GGML_OP_IM2COL) { -+ const int32_t s0 = tensor->op_params[0]; -+ const int32_t s1 = tensor->op_params[1]; -+ const int32_t p0 = tensor->op_params[2]; -+ const int32_t p1 = tensor->op_params[3]; -+ const int32_t d0 = tensor->op_params[4]; -+ const int32_t d1 = tensor->op_params[5]; -+ -+ const bool is_2D = tensor->op_params[6] == 1; -+ tensor_clone = ggml_im2col(ggml_ctx, src0_clone, src1_clone, s0, s1, p0, p1, d0, d1, is_2D, tensor->type); -+ } else if (tensor->op == GGML_OP_TIMESTEP_EMBEDDING) { -+ const int32_t dim = tensor->op_params[0]; -+ const int32_t max_period = tensor->op_params[1]; -+ tensor_clone = ggml_timestep_embedding(ggml_ctx, src0_clone, dim, max_period); -+ } else if (tensor->op == GGML_OP_POOL_2D) { -+ enum ggml_op_pool op = static_cast(tensor->op_params[0]); -+ const int32_t k0 = tensor->op_params[1]; -+ const int32_t k1 = tensor->op_params[2]; -+ const int32_t s0 = tensor->op_params[3]; -+ const int32_t s1 = tensor->op_params[4]; -+ const int32_t p0 = tensor->op_params[5]; -+ const int32_t p1 = tensor->op_params[6]; -+ -+ tensor_clone = ggml_pool_2d(ggml_ctx, src0_clone, op, k0, k1, s0, s1, p0, p1); -+ } else if (tensor->op == GGML_OP_LEAKY_RELU) { -+ const float * op_params = (const float *)tensor->op_params; -+ tensor_clone = ggml_leaky_relu(ggml_ctx, src0_clone, op_params[0], false); -+ } else if (tensor->op == GGML_OP_RWKV_WKV6) { -+ tensor_clone = ggml_rwkv_wkv6(ggml_ctx, tensor->src[0], tensor->src[1], tensor->src[2], tensor->src[3], -+ tensor->src[4], tensor->src[5]); -+ } -+ else { -+ std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl; -+ GGML_ABORT("fatal error"); -+ } -+ -+ ggml_cgraph * cgraph = ggml_new_graph(ggml_ctx); -+ ggml_build_forward_expand(cgraph, tensor_clone); -+ -+ ggml_graph_compute_with_ctx(ggml_ctx, cgraph, 8); -+ -+ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { -+ ggml_vk_print_tensor(tensor_clone, "tensor_clone"); -+ } -+ -+ comp_size = ggml_nbytes(tensor_clone); -+ -+ comp_result = malloc(comp_size); -+ memcpy(comp_result, tensor_clone->data, comp_size); -+ memcpy(comp_nb, tensor_clone->nb, sizeof(size_t) * GGML_MAX_DIMS); -+ -+ if (src0 != nullptr) { -+ free(src0_buffer); -+ } -+ if (src1 != nullptr) { -+ free(src1_buffer); -+ } -+ -+ ggml_free(ggml_ctx); -+ -+ VK_LOG_DEBUG("END ggml_vk_check_results_0(" << tensor->name << ")"); -+} -+ -+static void ggml_vk_check_results_1(ggml_tensor * tensor) { -+ if (tensor->op == GGML_OP_TRANSPOSE) { -+ return; -+ } -+ if (!(vk_output_tensor > 0 && vk_output_tensor == check_counter) && check_counter <= vk_skip_checks) { -+ return; -+ } -+ -+ VK_LOG_DEBUG("ggml_vk_check_results_1(" << tensor->name << ")"); -+ -+ ggml_tensor * src0 = tensor->src[0]; -+ ggml_tensor * src1 = tensor->src[1]; -+ ggml_tensor * src2 = tensor->src[2]; -+ -+ void * tensor_data = tensor->data; -+ -+ if (ggml_backend_buffer_is_vk(tensor->buffer)) { -+ size_t tensor_size = ggml_nbytes(tensor); -+ tensor_data = malloc(tensor_size); -+ -+ ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context; -+ -+ vk_buffer& buffer_gpu = buf_ctx->dev_buffer; -+ uint64_t offset = vk_tensor_offset(tensor) + tensor->view_offs; -+ if (offset + tensor_size >= buffer_gpu->size) { -+ tensor_size = buffer_gpu->size - offset; -+ } -+ -+ ggml_vk_buffer_read(buffer_gpu, offset, tensor_data, tensor_size); -+ } -+ -+ float first_error_result = -1.0f; -+ float first_error_correct = -1.0f; -+ std::array first_error = { -1, -1, -1, -1 }; -+ double avg_err = 0.0; -+ size_t counter = 0; -+ -+ for (int i3 = 0; i3 < tensor->ne[3]; i3++) { -+ for (int i2 = 0; i2 < tensor->ne[2]; i2++) { -+ for (int i1 = 0; i1 < tensor->ne[1]; i1++) { -+ for (int i0 = 0; i0 < tensor->ne[0]; i0++) { -+ const bool buffer_size_fit = i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0] < comp_size; -+ float correct = 0.0f; -+ float result = 0.0f; -+ -+ if (buffer_size_fit) { -+ if (tensor->type == GGML_TYPE_F32) { -+ correct = *(float *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0]); -+ result = *(float *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0]); -+ } else if (tensor->type == GGML_TYPE_F16) { -+ correct = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0])); -+ result = ggml_fp16_to_fp32(*(ggml_fp16_t *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0])); -+ } else if (tensor->type == GGML_TYPE_I32) { -+ correct = *(int32_t *) ((char *) comp_result + i3*comp_nb[3] + i2*comp_nb[2] + i1*comp_nb[1] + i0*comp_nb[0]); -+ result = *(int32_t *) ((char *) tensor_data + i3*tensor->nb[3] + i2*tensor->nb[2] + i1*tensor->nb[1] + i0*tensor->nb[0]); -+ } else { -+ std::cerr << "Results check not implemented for type " << ggml_type_name(tensor->type) << std::endl; -+ } -+ } else { -+ std::cerr << "Missing debug code for type " << ggml_type_name(tensor->type) << std::endl; -+ GGML_ABORT("fatal error"); -+ } -+ -+ if ((std::isnan(correct) != std::isnan(result)) || (std::isinf(correct) != std::isinf(result)) || !buffer_size_fit) { -+ std::cerr << "ERROR: Invalid value in " << ggml_op_name(tensor->op) << " i3=" << i3 << " i2=" << i2 << " i1=" << i1 << " i0=" << i0 << " result=" << result << " correct=" << correct << " avg_err=" << (avg_err / counter) << std::endl; -+ std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; -+ if (src0 != nullptr) { -+ std::cerr << "src0=" << src0 << " src0->name=" << src0->name << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; -+ } -+ if (src1 != nullptr) { -+ std::cerr << "src1=" << src1 << " src1->name=" << src1->name << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; -+ } -+ if (src2 != nullptr) { -+ std::cerr << "src2=" << src2 << " src2->name=" << src2->name << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; -+ } -+ std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; -+ std::cerr << std::endl << "Result:" << std::endl; -+ ggml_vk_print_tensor_area(tensor, tensor_data, i0, i1, i2, i3); -+ std::cerr << std::endl << "Correct:" << std::endl; -+ ggml_vk_print_tensor_area(tensor, comp_result, i0, i1, i2, i3); -+ std::cerr << std::endl; -+ std::vector done; -+ ggml_vk_print_graph_origin(tensor, done); -+ GGML_ABORT("fatal error"); -+ } -+ if (first_error[0] == -1 && std::fabs(correct - result) > 0.1f) { -+ first_error[0] = i0; -+ first_error[1] = i1; -+ first_error[2] = i2; -+ first_error[3] = i3; -+ first_error_result = result; -+ first_error_correct = correct; -+ } -+ -+ // Special case, value is infinite, avoid NaN result in avg_err -+ // NaN also appears in results, if both are nan error is 0 -+ if (!std::isinf(correct) && !std::isinf(result) && !std::isnan(correct) && !std::isnan(result)) { -+ avg_err += std::fabs(correct - result); -+ } -+ counter++; -+ } -+ } -+ } -+ } -+ -+ avg_err /= counter; -+ -+ if (vk_output_tensor > 0 && vk_output_tensor == check_counter) { -+ std::cerr << "TENSOR CHECK: avg_err=" << avg_err << " in " << ggml_op_name(tensor->op) << " (check " << check_counter << ")" << std::endl; -+ std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; -+ if (src0 != nullptr) { -+ std::cerr << "src0=" << src0 << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; -+ } -+ if (src1 != nullptr) { -+ std::cerr << "src1=" << src1 << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; -+ } -+ if (src2 != nullptr) { -+ std::cerr << "src2=" << src2 << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; -+ } -+ std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; -+ std::cerr << std::endl << "Result:" << std::endl; -+ ggml_vk_print_tensor_area(tensor, tensor_data, 5, 5, 0, 0); -+ std::cerr << std::endl << "Correct:" << std::endl; -+ ggml_vk_print_tensor_area(tensor, comp_result, 5, 5, 0, 0); -+ std::cerr << std::endl; -+ std::vector done; -+ ggml_vk_print_graph_origin(tensor, done); -+ } -+ -+ if (avg_err > 0.05 || std::isnan(avg_err)) { -+ std::cerr << "ERROR: avg_err=" << avg_err << " in " << ggml_op_name(tensor->op) << " (check " << check_counter << ")" << std::endl; -+ std::cerr << "tensor=" << tensor << " tensor->name=" << tensor->name << " tensor->type: " << ggml_type_name(tensor->type) << " ne0=" << tensor->ne[0] << " nb0=" << tensor->nb[0] << " ne1=" << tensor->ne[1] << " nb1=" << tensor->nb[1] << " ne2=" << tensor->ne[2] << " nb2=" << tensor->nb[2] << " ne3=" << tensor->ne[3] << " nb3=" << tensor->nb[3] << " offset=" << tensor->view_offs << std::endl; -+ if (src0 != nullptr) { -+ std::cerr << "src0=" << src0 << " op=" << ggml_op_name(src0->op) << " type=" << ggml_type_name(src0->type) << " ne0=" << src0->ne[0] << " nb0=" << src0->nb[0] << " ne1=" << src0->ne[1] << " nb1=" << src0->nb[1] << " ne2=" << src0->ne[2] << " nb2=" << src0->nb[2] << " ne3=" << src0->ne[3] << " nb3=" << src0->nb[3] << " offset=" << src0->view_offs << std::endl; -+ } -+ if (src1 != nullptr) { -+ std::cerr << "src1=" << src1 << " op=" << ggml_op_name(src1->op) << " type=" << ggml_type_name(src1->type) << " ne0=" << src1->ne[0] << " nb0=" << src1->nb[0] << " ne1=" << src1->ne[1] << " nb1=" << src1->nb[1] << " ne2=" << src1->ne[2] << " nb2=" << src1->nb[2] << " ne3=" << src1->ne[3] << " nb3=" << src1->nb[3] << " offset=" << src1->view_offs << std::endl; -+ } -+ if (src2 != nullptr) { -+ std::cerr << "src2=" << src2 << " op=" << ggml_op_name(src2->op) << " type=" << ggml_type_name(src2->type) << " ne0=" << src2->ne[0] << " nb0=" << src2->nb[0] << " ne1=" << src2->ne[1] << " nb1=" << src2->nb[1] << " ne2=" << src2->ne[2] << " nb2=" << src2->nb[2] << " ne3=" << src2->ne[3] << " nb3=" << src2->nb[3] << " offset=" << src2->view_offs << std::endl; -+ } -+ std::cerr << "First error: result=" << first_error_result << " correct=" << first_error_correct << " i3=" << first_error[3] << " i2=" << first_error[2] << " i1=" << first_error[1] << " i0=" << first_error[0] << std::endl; -+ std::cerr << std::endl << "Result:" << std::endl; -+ ggml_vk_print_tensor_area(tensor, tensor_data, first_error[0], first_error[1], first_error[2], first_error[3]); -+ std::cerr << std::endl << "Correct:" << std::endl; -+ ggml_vk_print_tensor_area(tensor, comp_result, first_error[0], first_error[1], first_error[2], first_error[3]); -+ std::cerr << std::endl; -+ std::vector done; -+ ggml_vk_print_graph_origin(tensor, done); -+ GGML_ABORT("fatal error"); -+ } else { -+ std::cerr << check_counter << " " << tensor->name << " op=" << ggml_op_name(tensor->op) << " avg_err=" << avg_err << std::endl; -+ } -+ -+ free(comp_result); -+ comp_result = nullptr; -+ comp_size = 0; -+ -+ if (ggml_backend_buffer_is_vk(tensor->buffer)) { -+ free(tensor_data); -+ } -+ -+ VK_LOG_DEBUG("END ggml_vk_check_results_1(" << tensor->name << ")"); -+} -+#endif -+ -+GGML_BACKEND_DL_IMPL(ggml_backend_vk_reg) -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt -new file mode 100644 -index 00000000..bd0c74cb ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt -@@ -0,0 +1,9 @@ -+find_package (Threads REQUIRED) -+find_package(Vulkan COMPONENTS glslc REQUIRED) -+ -+set(TARGET vulkan-shaders-gen) -+add_executable(${TARGET} vulkan-shaders-gen.cpp) -+install(TARGETS ${TARGET} RUNTIME) -+target_compile_features(${TARGET} PRIVATE cxx_std_17) -+target_link_libraries(vulkan-shaders-gen PUBLIC Threads::Threads) -+target_link_libraries(vulkan-shaders-gen PRIVATE Vulkan::Vulkan) -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp -new file mode 100644 -index 00000000..d896f1ef ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/acc.comp -@@ -0,0 +1,29 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_binary_head.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ const uint idx = gl_GlobalInvocationID.x; -+ if (idx >= p.ne) { -+ return; -+ } -+ -+ const uint offset = p.param3; -+ const uint src1_i = idx - offset; -+ const uint oz = src1_i / p.nb02; -+ const uint oy = (src1_i - (oz * p.nb02)) / p.nb01; -+ const uint ox = src1_i % p.nb01; -+ -+ uint i00, i01, i02, i03; -+ get_indices(idx, i00, i01, i02, i03); -+ -+ if (ox < p.ne10 && oy < p.ne11 && oz < p.ne12) { -+ data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) + FLOAT_TYPE(data_b[get_boffset() + ox + oy * p.ne10 + oz * p.ne10 * p.ne11])); -+ } else { -+ data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)])); -+ } -+} -+ -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp -new file mode 100644 -index 00000000..2b4085c4 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/add.comp -@@ -0,0 +1,29 @@ -+#version 450 -+ -+#extension GL_EXT_shader_16bit_storage : require -+ -+#include "types.comp" -+#include "generic_binary_head.comp" -+ -+const uint num_threads = 256; -+ -+layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ uint idx = get_idx(); -+ -+ // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation -+ const uint num_iter = 2; -+ -+ [[unroll]] for (uint i = 0; i < num_iter; ++i) { -+ if (idx >= p.ne) { -+ continue; -+ } -+ uint i00, i01, i02, i03; -+ get_indices(idx, i00, i01, i02, i03); -+ -+ data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) + FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); -+ -+ idx += num_threads; -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp -new file mode 100644 -index 00000000..d4fa45b1 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp -@@ -0,0 +1,69 @@ -+#version 450 -+ -+#include "types.comp" -+ -+#define BLOCK_SIZE 1024 -+#define ASC 0 -+ -+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) buffer D {int data_d[];}; -+ -+layout (push_constant) uniform parameter { -+ uint ncols; -+ uint ncols_pad; -+ uint order; -+} p; -+ -+shared int dst_row[BLOCK_SIZE]; -+ -+void swap(uint idx0, uint idx1) { -+ int tmp = dst_row[idx0]; -+ dst_row[idx0] = dst_row[idx1]; -+ dst_row[idx1] = tmp; -+} -+ -+void main() { -+ // bitonic sort -+ const int col = int(gl_LocalInvocationID.x); -+ const uint row = gl_WorkGroupID.y; -+ -+ const uint row_offset = row * p.ncols; -+ -+ // initialize indices -+ if (col < p.ncols_pad) { -+ dst_row[col] = col; -+ } -+ barrier(); -+ -+ for (uint k = 2; k <= p.ncols_pad; k *= 2) { -+ for (uint j = k / 2; j > 0; j /= 2) { -+ const uint ixj = col ^ j; -+ if (col < p.ncols_pad && ixj > col) { -+ if ((col & k) == 0) { -+ if (dst_row[col] >= p.ncols || -+ (dst_row[ixj] < p.ncols && (p.order == ASC ? -+ data_a[row_offset + dst_row[col]] > data_a[row_offset + dst_row[ixj]] : -+ data_a[row_offset + dst_row[col]] < data_a[row_offset + dst_row[ixj]])) -+ ) { -+ swap(col, ixj); -+ } -+ } else { -+ if (dst_row[ixj] >= p.ncols || -+ (dst_row[col] < p.ncols && (p.order == ASC ? -+ data_a[row_offset + dst_row[col]] < data_a[row_offset + dst_row[ixj]] : -+ data_a[row_offset + dst_row[col]] > data_a[row_offset + dst_row[ixj]])) -+ ) { -+ swap(col, ixj); -+ } -+ } -+ } -+ barrier(); -+ } -+ } -+ -+ if (col < p.ncols) { -+ data_d[row_offset + col] = dst_row[col]; -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp -new file mode 100644 -index 00000000..1e5cb8da ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/clamp.comp -@@ -0,0 +1,17 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_unary_head.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ const uint idx = get_idx(); -+ -+ if (idx >= p.ne) { -+ return; -+ } -+ -+ const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); -+ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(val < p.param1 ? p.param1 : (val > p.param2 ? p.param2 : val)); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp -new file mode 100644 -index 00000000..9ee2f1fa ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/concat.comp -@@ -0,0 +1,41 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_binary_head.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; -+ const int dim = p.param3; -+ -+ if (idx >= p.ne) { -+ return; -+ } -+ -+ const uint i3 = idx / (p.ne22*p.ne21*p.ne20); -+ const uint i3_offset = i3 * p.ne22*p.ne21*p.ne20; -+ const uint i2 = (idx - i3_offset) / (p.ne21*p.ne20); -+ const uint i2_offset = i2*p.ne21*p.ne20; -+ const uint i1 = (idx - i3_offset - i2_offset) / p.ne20; -+ const uint i0 = idx - i3_offset - i2_offset - i1*p.ne20; -+ -+ uint o[4] = {0, 0, 0, 0}; -+ o[dim] = dim == 0 ? p.ne00 : (dim == 1 ? p.ne01 : (dim == 2 ? p.ne02 : p.ne03)); -+ -+ const uint src0_idx = i3*p.nb03 + i2*p.nb02 + i1*p.nb01 + i0*p.nb00; -+ const uint src1_idx = (i3 - o[3])*p.nb13 + (i2 - o[2])*p.nb12 + (i1 - o[1])*p.nb11 + (i0 - o[0])*p.nb10; -+ const uint dst_idx = i3*p.nb23 + i2*p.nb22 + i1*p.nb21 + i0*p.nb20; -+ -+ const bool is_src0 = i0 < p.ne00 && i1 < p.ne01 && i2 < p.ne02 && i3 < p.ne03; -+ -+#ifndef OPTIMIZATION_ERROR_WORKAROUND -+ data_d[get_doffset() + dst_idx] = D_TYPE(is_src0 ? data_a[get_aoffset() + src0_idx] : data_b[get_boffset() + src1_idx]); -+#else -+ if (is_src0) { -+ data_d[get_doffset() + dst_idx] = data_a[get_aoffset() + src0_idx]; -+ } else { -+ data_d[get_doffset() + dst_idx] = data_b[get_boffset() + src1_idx]; -+ } -+#endif -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp -new file mode 100644 -index 00000000..dd828c23 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp -@@ -0,0 +1,42 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_unary_head.comp" -+ -+#extension GL_EXT_control_flow_attributes : require -+ -+const uint num_threads = 128; -+ -+layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ uint idx = get_idx(); -+ -+ // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation -+ const uint num_iter = 4; -+ -+ // fast path for when all four iterations are in-bounds -+ if (idx + (num_iter-1)*num_threads < p.ne) { -+ [[unroll]] for (uint i = 0; i < num_iter; ++i) { -+#ifndef OPTIMIZATION_ERROR_WORKAROUND -+ data_d[get_doffset() + idx] = D_TYPE(data_a[get_aoffset() + idx]); -+#else -+ data_d[get_doffset() + idx] = data_a[get_aoffset() + idx]; -+#endif -+ idx += num_threads; -+ } -+ } else { -+ [[unroll]] for (uint i = 0; i < num_iter; ++i) { -+ if (idx >= p.ne) { -+ continue; -+ } -+ -+#ifndef OPTIMIZATION_ERROR_WORKAROUND -+ data_d[get_doffset() + idx] = D_TYPE(data_a[get_aoffset() + idx]); -+#else -+ data_d[get_doffset() + idx] = data_a[get_aoffset() + idx]; -+#endif -+ idx += num_threads; -+ } -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp -new file mode 100644 -index 00000000..29c90649 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp -@@ -0,0 +1,20 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_unary_head.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ const uint idx = get_idx(); -+ -+ if (idx >= p.ne) { -+ return; -+ } -+ -+#ifndef OPTIMIZATION_ERROR_WORKAROUND -+ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(data_a[get_aoffset() + src0_idx(idx)]); -+#else -+ data_d[get_doffset() + dst_idx(idx)] = data_a[get_aoffset() + src0_idx(idx)]; -+#endif -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp -new file mode 100644 -index 00000000..0b8d02f5 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/cos.comp -@@ -0,0 +1,17 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_unary_head.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ const uint idx = get_idx(); -+ -+ if (idx >= p.ne) { -+ return; -+ } -+ -+ const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); -+ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(cos(val)); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp -new file mode 100644 -index 00000000..a4d3fca5 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f32.comp -@@ -0,0 +1,20 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {float data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ const uint i = gl_GlobalInvocationID.x * 16; -+ -+ if (i >= p.nel) { -+ return; -+ } -+ -+ [[unroll]] for (uint l = 0; l < 16; l++) { -+ data_b[i + l] = D_TYPE(data_a[i + l]); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp -new file mode 100644 -index 00000000..91bb8f8d ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp -@@ -0,0 +1,118 @@ -+#if !defined(DATA_A_F32) && !defined(DATA_A_F16) -+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -+#endif -+ -+#include "types.comp" -+ -+#if defined(A_TYPE_PACKED16) -+layout (binding = 0) readonly buffer A_PACKED16 {A_TYPE_PACKED16 data_a_packed16[];}; -+#endif -+#if defined(A_TYPE_PACKED32) -+layout (binding = 0) readonly buffer A_PACKED32 {A_TYPE_PACKED32 data_a_packed32[];}; -+#endif -+ -+#if defined(DATA_A_F32) -+vec2 dequantize(uint ib, uint iqs, uint a_offset) { -+ return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]); -+} -+#endif -+ -+#if defined(DATA_A_F16) -+vec2 dequantize(uint ib, uint iqs, uint a_offset) { -+ return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]); -+} -+#endif -+ -+#if defined(DATA_A_Q4_0) -+vec2 dequantize(uint ib, uint iqs, uint a_offset) { -+ const uint vui = uint(data_a[a_offset + ib].qs[iqs]); -+ return (vec2(vui & 0xF, vui >> 4) - 8.0f); -+} -+vec4 dequantize4(uint ib, uint iqs, uint a_offset) { -+ const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); -+ return (vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, vui >> 12) - 8.0f); -+} -+#endif -+ -+#if defined(DATA_A_Q4_1) -+vec2 dequantize(uint ib, uint iqs, uint a_offset) { -+ const uint vui = uint(data_a[a_offset + ib].qs[iqs]); -+ return vec2(vui & 0xF, vui >> 4); -+} -+vec4 dequantize4(uint ib, uint iqs, uint a_offset) { -+ const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); -+ return vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, vui >> 12); -+} -+#endif -+ -+#if defined(DATA_A_Q5_0) -+vec2 dequantize(uint ib, uint iqs, uint a_offset) { -+ const uint uint_qh = uint(data_a[a_offset + ib].qh[1]) << 16 | data_a[a_offset + ib].qh[0]; -+ const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); -+ const uint vui = uint(data_a[a_offset + ib].qs[iqs]); -+ return (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f); -+} -+vec4 dequantize4(uint ib, uint iqs, uint a_offset) { -+ const uint uint_qh = uint(data_a_packed16[a_offset + ib].qh[1]) << 16 | data_a_packed16[a_offset + ib].qh[0]; -+ const ivec2 qh0 = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); -+ const ivec2 qh1 = ivec2(((uint_qh >> (iqs + 1)) << 4) & 0x10, (uint_qh >> (iqs + 13)) & 0x10); -+ const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); -+ return (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f); -+} -+#endif -+ -+#if defined(DATA_A_Q5_1) -+vec2 dequantize(uint ib, uint iqs, uint a_offset) { -+ const uint uint_qh = data_a[a_offset + ib].qh; -+ const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); -+ const uint vui = uint(data_a[a_offset + ib].qs[iqs]); -+ return vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y); -+} -+vec4 dequantize4(uint ib, uint iqs, uint a_offset) { -+ const uint uint_qh = data_a_packed16[a_offset + ib].qh; -+ const ivec2 qh0 = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); -+ const ivec2 qh1 = ivec2(((uint_qh >> (iqs + 1)) << 4) & 0x10, (uint_qh >> (iqs + 13)) & 0x10); -+ const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); -+ return vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y); -+} -+#endif -+ -+#if defined(DATA_A_Q8_0) -+vec2 dequantize(uint ib, uint iqs, uint a_offset) { -+ return vec2(int(data_a[a_offset + ib].qs[iqs]), int(data_a[a_offset + ib].qs[iqs + 1])); -+} -+vec4 dequantize4(uint ib, uint iqs, uint a_offset) { -+ uint32_t v0 = data_a_packed16[a_offset + ib].qs[iqs/2]; -+ uint32_t v1 = data_a_packed16[a_offset + ib].qs[iqs/2 + 1]; -+ return vec4(int8_t(v0 & 0xFF), int8_t(v0 >> 8), int8_t(v1 & 0xFF), int8_t(v1 >> 8)); -+} -+#endif -+ -+#if defined(DATA_A_IQ4_NL) -+vec2 dequantize(uint ib, uint iqs, uint a_offset) { -+ const uint vui = uint(data_a[a_offset + ib].qs[iqs]); -+ return vec2(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[vui >> 4]); -+} -+vec4 dequantize4(uint ib, uint iqs, uint a_offset) { -+ const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]); -+ return vec4(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[(vui >> 4) & 0xF], kvalues_iq4nl[(vui >> 8) & 0xF], kvalues_iq4nl[vui >> 12]); -+} -+#endif -+ -+#if defined(DATA_A_F32) || defined(DATA_A_F16) -+vec2 get_dm(uint ib, uint a_offset) { -+ return vec2(0, 0); -+} -+#endif -+ -+#if defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ4_NL) -+vec2 get_dm(uint ib, uint a_offset) { -+ return vec2(float(data_a[a_offset + ib].d), 0); -+} -+#endif -+ -+#if defined(DATA_A_Q4_1) || defined(DATA_A_Q5_1) -+vec2 get_dm(uint ib, uint a_offset) { -+ return vec2(float(data_a[a_offset + ib].d), float(data_a[a_offset + ib].m)); -+} -+#endif -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp -new file mode 100644 -index 00000000..94b78598 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp -@@ -0,0 +1,325 @@ -+ -+#include "types.comp" -+ -+layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ4_0 { -+ block_q4_0_packed16 block; -+}; -+ -+float16_t dequantFuncQ4_0(const in decodeBufQ4_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ const float16_t d = bl.block.d; -+ const uint idx = coordInBlock[1]; -+ const uint shift = (idx & 0x10) >> 2; -+ uint32_t qs = uint32_t(bl.block.qs[(idx & 0xE) >> 1]); -+ qs >>= shift; -+ qs &= 0x0F0F; -+ qs = unpack8(qs)[idx & 1]; -+ float16_t ret = (float16_t(qs) - float16_t(8)) * d; -+ return ret; -+} -+ -+layout(buffer_reference, std430, buffer_reference_align = 4) buffer decodeBufQ4_1 { -+ block_q4_1 block; -+}; -+ -+float16_t dequantFuncQ4_1(const in decodeBufQ4_1 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ const float16_t d = bl.block.d; -+ const float16_t m = bl.block.m; -+ const uint idx = coordInBlock[1]; -+ const uint iqs = idx & 0xF; -+ const uint shift = (idx & 0x10) >> 2; -+ uint32_t qs = bl.block.qs[iqs]; -+ qs >>= shift; -+ qs &= 0xF; -+ float16_t ret = float16_t(qs) * d + m; -+ return ret; -+} -+ -+layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ5_0 { -+ block_q5_0 block; -+}; -+ -+float16_t dequantFuncQ5_0(const in decodeBufQ5_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ const float16_t d = bl.block.d; -+ const uint idx = coordInBlock[1]; -+ const uint iqs = idx & 0xF; -+ -+ const uint uint_qh = uint(bl.block.qh[1]) << 16 | bl.block.qh[0]; -+ const uint qh = ((uint_qh >> idx) << 4) & 0x10; -+ -+ const uint shift = (idx & 0x10) >> 2; -+ uint32_t qs = bl.block.qs[iqs]; -+ qs >>= shift; -+ qs &= 0xF; -+ -+ float16_t ret = (float16_t(qs | qh) - float16_t(16)) * d; -+ return ret; -+} -+ -+layout(buffer_reference, std430, buffer_reference_align = 8) buffer decodeBufQ5_1 { -+ block_q5_1 block; -+}; -+ -+float16_t dequantFuncQ5_1(const in decodeBufQ5_1 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ const float16_t d = bl.block.d; -+ const float16_t m = bl.block.m; -+ const uint idx = coordInBlock[1]; -+ const uint iqs = idx & 0xF; -+ -+ const uint uint_qh = bl.block.qh; -+ const uint qh = ((uint_qh >> idx) << 4) & 0x10; -+ -+ const uint shift = (idx & 0x10) >> 2; -+ uint32_t qs = bl.block.qs[iqs]; -+ qs >>= shift; -+ qs &= 0xF; -+ -+ float16_t ret = float16_t(qs | qh) * d + m; -+ return ret; -+} -+ -+layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ8_0 { -+ block_q8_0_packed16 block; -+}; -+ -+float16_t dequantFuncQ8_0(const in decodeBufQ8_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ const float16_t d = bl.block.d; -+ const uint idx = coordInBlock[1]; -+ const uint iqs = idx; -+ -+ // Load 16b and select the byte for this element -+ int32_t qs = unpack8(int32_t(bl.block.qs[(iqs & 0x1E) >> 1]))[iqs & 1]; -+ float16_t ret = float16_t(qs) * d; -+ return ret; -+} -+ -+layout(buffer_reference, std430, buffer_reference_align = 4) buffer decodeBufQ2_K { -+ block_q2_K block; -+}; -+ -+float16_t dequantFuncQ2_K(const in decodeBufQ2_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ const f16vec2 d = bl.block.d; -+ const uint idx = coordInBlock[1]; -+ const uint iqs = idx; -+ -+ const uint qsi = (iqs / 128) * 32 + (iqs % 32); // 0..31 -+ const uint scalesi = iqs / 16; // 0..15 -+ const uint qsshift = ((iqs % 128) / 32) * 2; // 0,2,4,6 -+ -+ uint32_t qs = bl.block.qs[qsi]; -+ const uint scales = bl.block.scales[scalesi]; -+ float16_t ret = d.x * float16_t(scales & 0xF) * float16_t((qs >> qsshift) & 3) - d.y * float16_t(scales >> 4); -+ return ret; -+} -+ -+layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ3_K { -+ block_q3_K block; -+}; -+ -+float16_t dequantFuncQ3_K(const in decodeBufQ3_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ const uint idx = coordInBlock[1]; -+ const uint iqs = idx; -+ -+ const uint n = iqs / 128; // 0,1 -+ const uint qsi = n * 32 + (iqs % 32); // 0..63 -+ const uint hmi = (iqs % 32); // 0..31 -+ const uint j = (iqs % 128) / 8; // 0..15 -+ const uint is = iqs / 16; // 0..15 -+ const uint halfsplit = ((iqs % 128) / 32); // 0,1,2,3 -+ const uint qsshift = halfsplit * 2; // 0,2,4,6 -+ const uint m = 1 << (4 * n + halfsplit); // 1,2,4,8,16,32,64,128 -+ -+ uint32_t scaleidx0 = (is < 8) ? is : (is-8); -+ uint32_t scaleidx0shift = (is < 8) ? 0 : 4; -+ uint32_t scaleidx1 = is + 8 - (is/4)*4; -+ uint32_t scaleidx1shift = (is/4)*2; -+ -+ const int8_t us = int8_t(((bl.block.scales[scaleidx0] >> scaleidx0shift) & 0xF) | (((bl.block.scales[scaleidx1] >> scaleidx1shift) & 3) << 4)); -+ -+ const float16_t dl = bl.block.d * float16_t(us - 32); -+ -+ float16_t ret = dl * float16_t(int8_t((bl.block.qs[qsi ] >> qsshift) & 3) - (((bl.block.hmask[hmi ] & m) != 0) ? 0 : 4)); -+ -+ return ret; -+} -+ -+layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ4_K { -+ block_q4_K block; -+}; -+ -+layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ4_K_packed16 { -+ block_q4_K_packed16 block; -+}; -+ -+float16_t dequantFuncQ4_K(const in decodeBufQ4_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ decodeBufQ4_K_packed16 bl16 = decodeBufQ4_K_packed16(bl); -+ const uint idx = coordInBlock[1]; -+ -+ const uint b = (idx & 0x20) >> 5; // 0,1 -+ const uint is = (idx & 0xE0) >> 5; // 0..7 -+ -+ const f16vec2 loadd = bl.block.d; -+ -+ uint32_t sc; -+ uint32_t mbyte; -+ -+ uint32_t scidx0 = (is < 4) ? is : (is + 4); -+ uint32_t scidx1 = (is < 4) ? is : (is - 4); -+ uint32_t scidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ uint32_t scidxshift1 = (is < 4) ? 0 : 2; -+ uint32_t mbidx0 = is + 4; -+ uint32_t mbidx1 = (is < 4) ? is + 4 : is; -+ uint32_t mbidxmask0 = (is < 4) ? 0xF : 0xF0; -+ uint32_t mbidxshift0 = (is < 4) ? 0 : 4; -+ uint32_t mbidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ uint32_t mbidxshift1 = (is < 4) ? 0 : 2; -+ -+ sc = uint8_t((bl.block.scales[scidx0] & 0xF) | ((bl.block.scales[scidx1] & scidxmask1) >> scidxshift1)); -+ mbyte = uint8_t(((bl.block.scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((bl.block.scales[mbidx1] & mbidxmask1) >> mbidxshift1)); -+ -+ const float16_t d = loadd.x * float16_t(sc); -+ const float16_t m = loadd.y * float16_t(mbyte); -+ -+ uint qs = uint32_t(bl16.block.qs[((idx & 0xC0) >> 2) + ((idx & 0x1E) >> 1)]); -+ qs = (qs >> (b * 4)) & 0x0F0F; -+ qs = unpack8(qs)[idx & 1]; -+ -+ float16_t ret = d * float16_t(qs) - m; -+ -+ return ret; -+} -+ -+layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ5_K { -+ block_q5_K block; -+}; -+ -+layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ5_K_packed16 { -+ block_q5_K_packed16 block; -+}; -+ -+float16_t dequantFuncQ5_K(const in decodeBufQ5_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ decodeBufQ5_K_packed16 bl16 = decodeBufQ5_K_packed16(bl); -+ const uint idx = coordInBlock[1]; -+ -+ const uint b = (idx & 0x20) >> 5; // 0,1 -+ const uint is = (idx & 0xE0) >> 5; // 0..7 -+ -+ const uint32_t hm = 0x0101 << is; -+ -+ const f16vec2 loadd = bl.block.d; -+ -+ uint32_t sc; -+ uint32_t mbyte; -+ -+ uint32_t scidx0 = (is < 4) ? is : (is + 4); -+ uint32_t scidx1 = (is < 4) ? is : (is - 4); -+ uint32_t scidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ uint32_t scidxshift1 = (is < 4) ? 0 : 2; -+ uint32_t mbidx0 = is + 4; -+ uint32_t mbidx1 = (is < 4) ? is + 4 : is; -+ uint32_t mbidxmask0 = (is < 4) ? 0xF : 0xF0; -+ uint32_t mbidxshift0 = (is < 4) ? 0 : 4; -+ uint32_t mbidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ uint32_t mbidxshift1 = (is < 4) ? 0 : 2; -+ -+ sc = uint8_t((bl.block.scales[scidx0] & 0xF) | ((bl.block.scales[scidx1] & scidxmask1) >> scidxshift1)); -+ mbyte = uint8_t(((bl.block.scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((bl.block.scales[mbidx1] & mbidxmask1) >> mbidxshift1)); -+ -+ const float16_t d = loadd.x * float16_t(sc); -+ const float16_t m = loadd.y * float16_t(mbyte); -+ -+ uint qh = uint32_t(bl16.block.qh[(idx & 0x1E) >> 1]); -+ qh = qh & hm; -+ qh = unpack8(qh)[idx & 1]; -+ -+ uint qs = uint32_t(bl16.block.qs[((idx & 0xC0) >> 2) + ((idx & 0x1E) >> 1)]); -+ qs = (qs >> (b * 4)) & 0x0F0F; -+ qs = unpack8(qs)[idx & 1]; -+ -+ float16_t ret = d * (float16_t(qs) + (qh != 0 ? float16_t(16) : float16_t(0))) - m; -+ -+ return ret; -+} -+ -+layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ6_K { -+ block_q6_K block; -+}; -+ -+layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ6_K_packed16 { -+ block_q6_K_packed16 block; -+}; -+ -+float16_t dequantFuncQ6_K(const in decodeBufQ6_K bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ decodeBufQ6_K_packed16 bl16 = decodeBufQ6_K_packed16(bl); -+ const uint idx = coordInBlock[1]; -+ -+ const uint b = (idx & 0x40) >> 6; // 0,1 -+ const uint qhshift = (idx & 0x60) >> 4; // 0,2,4,6 -+ const uint is = (idx & 0xF0) >> 4; // 0..15 -+ -+ const float16_t dscale = bl.block.d * float16_t(bl.block.scales[is]); -+ -+ uint ql = uint32_t(bl16.block.ql[((idx & 0x80) >> 2) + ((idx & 0x3E) >> 1)]); -+ ql = (ql >> (b * 4)) & 0x0F0F; -+ -+ uint qh = uint32_t(bl16.block.qh[((idx & 0x80) >> 3) + ((idx & 0x1E) >> 1)]); -+ qh = ((qh >> qhshift) & 0x0303) << 4; -+ -+ int q = unpack8(ql | qh)[idx & 1]; -+ -+ float16_t ret = dscale * float16_t(q - 32); -+ -+ return ret; -+} -+ -+#if defined(DATA_A_IQ4_NL) -+layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufIQ4_NL { -+ block_iq4_nl block; -+}; -+ -+float16_t dequantFuncIQ4_NL(const in decodeBufIQ4_NL bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ const float16_t d = bl.block.d; -+ const uint idx = coordInBlock[1]; -+ const uint iqs = idx & 0xF; -+ const uint shift = (idx & 0x10) >> 2; -+ uint32_t qs = bl.block.qs[iqs]; -+ qs >>= shift; -+ qs &= 0xF; -+ float16_t ret = float16_t(kvalues_iq4nl[qs]) * d; -+ return ret; -+} -+#endif -+ -+#if defined(DATA_A_Q4_0) -+#define dequantFuncA dequantFuncQ4_0 -+#elif defined(DATA_A_Q4_1) -+#define dequantFuncA dequantFuncQ4_1 -+#elif defined(DATA_A_Q5_0) -+#define dequantFuncA dequantFuncQ5_0 -+#elif defined(DATA_A_Q5_1) -+#define dequantFuncA dequantFuncQ5_1 -+#elif defined(DATA_A_Q8_0) -+#define dequantFuncA dequantFuncQ8_0 -+#elif defined(DATA_A_Q2_K) -+#define dequantFuncA dequantFuncQ2_K -+#elif defined(DATA_A_Q3_K) -+#define dequantFuncA dequantFuncQ3_K -+#elif defined(DATA_A_Q4_K) -+#define dequantFuncA dequantFuncQ4_K -+#elif defined(DATA_A_Q5_K) -+#define dequantFuncA dequantFuncQ5_K -+#elif defined(DATA_A_Q6_K) -+#define dequantFuncA dequantFuncQ6_K -+#elif defined(DATA_A_IQ4_NL) -+#define dequantFuncA dequantFuncIQ4_NL -+#endif -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp -new file mode 100644 -index 00000000..8d806435 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_head.comp -@@ -0,0 +1,13 @@ -+#extension GL_EXT_control_flow_attributes : require -+#extension GL_EXT_shader_16bit_storage : require -+ -+layout (push_constant) uniform parameter -+{ -+ uint M; -+ uint K; -+ uint stride_a; -+ uint stride_b; -+ uint nel; -+} p; -+ -+#include "types.comp" -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp -new file mode 100644 -index 00000000..8de14fc0 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp -@@ -0,0 +1,32 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {block_iq4_nl data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; -+ -+ init_iq4nl_shmem(); -+ -+ const uint tid = gl_LocalInvocationID.x % 64; -+ const uint il = tid/32; -+ const uint ir = tid%32; -+ const uint ib = 32*i + ir; -+ if (ib >= p.nel / 32) { -+ return; -+ } -+ -+ const uint q_idx = 8*il; -+ const uint b_idx = 1024*i + 32*ir + q_idx; -+ -+ const float d = float(data_a[ib].d); -+ -+ [[unroll]] for (uint l = 0; l < 8; ++l) { -+ data_b[b_idx + l + 0] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] & 0xF]); -+ data_b[b_idx + l + 16] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] >> 4]); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp -new file mode 100644 -index 00000000..157154af ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q2_k.comp -@@ -0,0 +1,34 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { -+ const uint i = gl_WorkGroupID.x * 256 + wgy; -+ if (i >= p.M * p.K / QUANT_K) { -+ return; -+ } -+ -+ const uint tid = gl_LocalInvocationID.x; -+ const uint ip = tid / 32; -+ const uint il = tid - 32 * ip; -+ const uint is = 8 * ip + il / 16; -+ -+ const uint y_idx = i * QUANT_K + 128 * ip + il; -+ -+ const uint ql_idx = 32 * ip + il; -+ const uint8_t qs = data_a[i].qs[32 * ip + il]; -+ -+ FLOAT_TYPE dall = FLOAT_TYPE(data_a[i].d.x); -+ FLOAT_TYPE dmin = FLOAT_TYPE(data_a[i].d.y); -+ data_b[y_idx + 0] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+0] & 0xF) * ((qs >> 0) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+0] >> 4)); -+ data_b[y_idx + 32] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+2] & 0xF) * ((qs >> 2) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+2] >> 4)); -+ data_b[y_idx + 64] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+4] & 0xF) * ((qs >> 4) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+4] >> 4)); -+ data_b[y_idx + 96] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+6] & 0xF) * ((qs >> 6) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+6] >> 4)); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp -new file mode 100644 -index 00000000..c17dd0d9 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q3_k.comp -@@ -0,0 +1,42 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { -+ const uint i = uint(gl_WorkGroupID.x * 256 + wgy); -+ if (i >= p.M * p.K / QUANT_K) { -+ return; -+ } -+ -+ const uint r = gl_LocalInvocationID.x / 4; -+ const uint tid = r / 2; -+ const uint is0 = r % 2; -+ const uint l0 = 16 * is0 + 4 * (gl_LocalInvocationID.x % 4); -+ const uint n = tid / 4; -+ const uint j = tid - 4*n; -+ -+ const uint8_t m = uint8_t(1 << (4*n + j)); -+ const uint is = 8*n + 2*j + is0; -+ const uint shift = 2*j; -+ -+ const int8_t us = int8_t(is < 4 ? (data_a[i].scales[is-0] & 0xF) | (((data_a[i].scales[is+8] >> 0) & 3) << 4) : -+ is < 8 ? (data_a[i].scales[is-0] & 0xF) | (((data_a[i].scales[is+4] >> 2) & 3) << 4) : -+ is < 12 ? (data_a[i].scales[is-8] >> 4) | (((data_a[i].scales[is+0] >> 4) & 3) << 4) : -+ (data_a[i].scales[is-8] >> 4) | (((data_a[i].scales[is-4] >> 6) & 3) << 4)); -+ const FLOAT_TYPE d_all = FLOAT_TYPE(data_a[i].d); -+ const FLOAT_TYPE dl = d_all * FLOAT_TYPE(us - 32); -+ -+ const uint y_idx = i * QUANT_K + 128 * n + 32 * j; -+ const uint qs_idx = 32*n; -+ -+ for (uint l = l0; l < l0 + 4; ++l) { -+ data_b[y_idx + l] = D_TYPE(dl * FLOAT_TYPE(int8_t((data_a[i].qs[qs_idx + l] >> shift) & 3) - (((data_a[i].hmask[l] & m) != 0) ? 0 : 4))); -+ } -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp -new file mode 100644 -index 00000000..40818532 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp -@@ -0,0 +1,30 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {block_q4_0 data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; -+ -+ const uint tid = gl_LocalInvocationID.x % 64; -+ const uint il = tid/32; -+ const uint ir = tid%32; -+ const uint ib = 32*i + ir; -+ if (ib >= p.nel / 32) { -+ return; -+ } -+ -+ const uint q_idx = 8*il; -+ const uint b_idx = 1024*i + 32*ir + q_idx; -+ -+ const float d = float(data_a[ib].d); -+ -+ [[unroll]] for (uint l = 0; l < 8; ++l) { -+ data_b[b_idx + l + 0] = D_TYPE(d * ((data_a[ib].qs[q_idx + l] & 0xF) - 8.0f)); -+ data_b[b_idx + l + 16] = D_TYPE(d * ((data_a[ib].qs[q_idx + l] >> 4) - 8.0f)); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp -new file mode 100644 -index 00000000..2f27eee6 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp -@@ -0,0 +1,32 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {block_q4_1 data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; -+ -+ const uint tid = gl_LocalInvocationID.x % 64; -+ const uint il = tid/32; -+ const uint ir = tid%32; -+ const uint ib = 32*i + ir; -+ if (ib >= p.nel / 32) { -+ return; -+ } -+ -+ const uint b_idx = 1024*i + 32*ir + 8*il; -+ -+ const float d = float(data_a[ib].d); -+ const float m = float(data_a[ib].m); -+ -+ const uint q_idx = 8*il; -+ -+ [[unroll]] for (uint l = 0; l < 8; ++l) { -+ data_b[b_idx + l + 0] = D_TYPE(d * (data_a[ib].qs[q_idx + l] & 0xF) + m); -+ data_b[b_idx + l + 16] = D_TYPE(d * (data_a[ib].qs[q_idx + l] >> 4) + m); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp -new file mode 100644 -index 00000000..987f113a ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_k.comp -@@ -0,0 +1,68 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { -+ const uint ib = gl_WorkGroupID.x * 256 + wgy; -+ if (ib >= p.M * p.K / QUANT_K) { -+ return; -+ } -+ -+ const uint tid = gl_LocalInvocationID.x; -+ const uint il = tid / 8; -+ const uint ir = tid % 8; -+ const uint is = 2 * il; -+ const uint n = 4; -+ -+ const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib].d.x); -+ const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib].d.y); -+ -+ const uint y_idx = ib * QUANT_K + 64 * il + n * ir; -+ const uint qs_idx = 32*il + n * ir; -+ -+ uint scidx0 = (is < 4) ? is : (is + 4); -+ uint scidx1 = (is < 4) ? is : (is - 4); -+ uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ uint scidxshift1 = (is < 4) ? 0 : 2; -+ uint mbidx0 = is + 4; -+ uint mbidx1 = (is < 4) ? is + 4 : is; -+ uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; -+ uint mbidxshift0 = (is < 4) ? 0 : 4; -+ uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ uint mbidxshift1 = (is < 4) ? 0 : 2; -+ -+ uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); -+ uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); -+ -+ const FLOAT_TYPE d1 = dall * sc; -+ const FLOAT_TYPE m1 = dmin * mbyte; -+ -+ scidx0 = (is < 4) ? is + 1 : (is + 5); -+ scidx1 = (is < 4) ? is + 1 : (is - 3); -+ scidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ scidxshift1 = (is < 4) ? 0 : 2; -+ mbidx0 = is + 5; -+ mbidx1 = (is < 4) ? is + 5 : is + 1; -+ mbidxmask0 = (is < 4) ? 0xF : 0xF0; -+ mbidxshift0 = (is < 4) ? 0 : 4; -+ mbidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ mbidxshift1 = (is < 4) ? 0 : 2; -+ -+ sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); -+ mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); -+ -+ const FLOAT_TYPE d2 = dall * sc; -+ const FLOAT_TYPE m2 = dmin * mbyte; -+ -+ [[unroll]] for (uint l = 0; l < n; ++l) { -+ data_b[y_idx + l ] = D_TYPE(d1 * FLOAT_TYPE(data_a[ib].qs[qs_idx + l] & 0xF) - m1); -+ data_b[y_idx + l + 32] = D_TYPE(d2 * FLOAT_TYPE(data_a[ib].qs[qs_idx + l] >> 4) - m2); -+ } -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp -new file mode 100644 -index 00000000..b20b8052 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp -@@ -0,0 +1,34 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {block_q5_0 data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; -+ -+ const uint tid = gl_LocalInvocationID.x % 64; -+ const uint il = tid/32; -+ const uint ir = tid%32; -+ const uint ib = 32*i + ir; -+ if (ib >= p.nel / 32) { -+ return; -+ } -+ -+ const uint b_idx = 1024*i + 32*ir + 8*il; -+ -+ const float d = float(data_a[ib].d); -+ const uint qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0]; -+ -+ const uint q_idx = 8*il; -+ -+ [[unroll]] for (uint l = 0; l < 8; ++l) { -+ const uint iqs = q_idx + l; -+ const uint vui = uint(data_a[ib].qs[iqs]); -+ data_b[b_idx + l + 0] = D_TYPE(d * (((vui & 0xF) | (((qh >> iqs) << 4) & 0x10)) - 16.0f)); -+ data_b[b_idx + l + 16] = D_TYPE(d * (((vui >> 4) | ((qh >> (iqs + 12)) & 0x10)) - 16.0f)); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp -new file mode 100644 -index 00000000..dc59fe3b ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp -@@ -0,0 +1,35 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {block_q5_1 data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; -+ -+ const uint tid = gl_LocalInvocationID.x % 64; -+ const uint il = tid/32; -+ const uint ir = tid%32; -+ const uint ib = 32*i + ir; -+ if (ib >= p.nel / 32) { -+ return; -+ } -+ -+ const uint b_idx = 1024*i + 32*ir + 8*il; -+ -+ const float d = float(data_a[ib].d); -+ const float m = float(data_a[ib].m); -+ const uint qh = data_a[ib].qh; -+ -+ const uint q_idx = 8*il; -+ -+ [[unroll]] for (uint l = 0; l < 8; ++l) { -+ const uint iqs = q_idx + l; -+ const uint vui = uint(data_a[ib].qs[iqs]); -+ data_b[b_idx + l + 0] = D_TYPE(d * (((vui & 0xF) | (((qh >> iqs) << 4) & 0x10))) + m); -+ data_b[b_idx + l + 16] = D_TYPE(d * (((vui >> 4) | ((qh >> (iqs + 12)) & 0x10))) + m); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp -new file mode 100644 -index 00000000..6db5403b ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_k.comp -@@ -0,0 +1,70 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { -+ const uint ib = gl_WorkGroupID.x * 256 + wgy; -+ if (ib >= p.M * p.K / QUANT_K) { -+ return; -+ } -+ -+ const uint tid = gl_LocalInvocationID.x; -+ const uint il = tid / 16; -+ const uint ir = tid % 16; -+ const uint is = 2 * il; -+ -+ const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib].d.x); -+ const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib].d.y); -+ -+ const uint y_idx = ib * QUANT_K + 64 * il + 2 * ir; -+ const uint qs_idx = 32*il + 2 * ir; -+ const uint qh_idx = 2 * ir; -+ -+ uint scidx0 = (is < 4) ? is : (is + 4); -+ uint scidx1 = (is < 4) ? is : (is - 4); -+ uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ uint scidxshift1 = (is < 4) ? 0 : 2; -+ uint mbidx0 = is + 4; -+ uint mbidx1 = (is < 4) ? is + 4 : is; -+ uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; -+ uint mbidxshift0 = (is < 4) ? 0 : 4; -+ uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ uint mbidxshift1 = (is < 4) ? 0 : 2; -+ -+ uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); -+ uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); -+ -+ const FLOAT_TYPE d1 = dall * sc; -+ const FLOAT_TYPE m1 = dmin * mbyte; -+ -+ scidx0 = (is < 4) ? is + 1 : (is + 5); -+ scidx1 = (is < 4) ? is + 1 : (is - 3); -+ scidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ scidxshift1 = (is < 4) ? 0 : 2; -+ mbidx0 = is + 5; -+ mbidx1 = (is < 4) ? is + 5 : is + 1; -+ mbidxmask0 = (is < 4) ? 0xF : 0xF0; -+ mbidxshift0 = (is < 4) ? 0 : 4; -+ mbidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ mbidxshift1 = (is < 4) ? 0 : 2; -+ -+ sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); -+ mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); -+ -+ const FLOAT_TYPE d2 = dall * sc; -+ const FLOAT_TYPE m2 = dmin * mbyte; -+ -+ const uint8_t hm1 = uint8_t(1 << (2 * il )); -+ const uint8_t hm2 = uint8_t(1 << (2 * il + 1)); -+ data_b[y_idx ] = D_TYPE(d1 * FLOAT_TYPE((data_a[ib].qs[qs_idx ] & 0xF) + (((data_a[ib].qh[qh_idx ] & hm1) != 0) ? 16 : 0)) - m1); -+ data_b[y_idx + 1] = D_TYPE(d1 * FLOAT_TYPE((data_a[ib].qs[qs_idx + 1] & 0xF) + (((data_a[ib].qh[qh_idx + 1] & hm1) != 0) ? 16 : 0)) - m1); -+ data_b[y_idx + 32] = D_TYPE(d2 * FLOAT_TYPE((data_a[ib].qs[qs_idx ] >> 4) + (((data_a[ib].qh[qh_idx ] & hm2) != 0) ? 16 : 0)) - m2); -+ data_b[y_idx + 33] = D_TYPE(d2 * FLOAT_TYPE((data_a[ib].qs[qs_idx + 1] >> 4) + (((data_a[ib].qh[qh_idx + 1] & hm2) != 0) ? 16 : 0)) - m2); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp -new file mode 100644 -index 00000000..0b913175 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q6_k.comp -@@ -0,0 +1,33 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { -+ const uint i = gl_WorkGroupID.x * 256 + wgy; -+ if (i >= p.M * p.K / QUANT_K) { -+ return; -+ } -+ const uint tid = gl_LocalInvocationID.x; -+ const uint ip = tid / 32; -+ const uint il = tid - 32 * ip; -+ const uint is = 8 * ip + il / 16; -+ -+ const uint y_idx = i * QUANT_K + 128 * ip + il; -+ -+ const uint ql_idx = 64 * ip + il; -+ const uint8_t qh = data_a[i].qh[32 * ip + il]; -+ -+ const FLOAT_TYPE d = FLOAT_TYPE(data_a[i].d); -+ -+ data_b[y_idx + 0] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 0] * (int8_t((data_a[i].ql[ql_idx + 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32))); -+ data_b[y_idx + 32] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 2] * (int8_t((data_a[i].ql[ql_idx + 32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32))); -+ data_b[y_idx + 64] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 4] * (int8_t((data_a[i].ql[ql_idx + 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32))); -+ data_b[y_idx + 96] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 6] * (int8_t((data_a[i].ql[ql_idx + 32] >> 4) | (((qh >> 6) & 3) << 4)) - 32))); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp -new file mode 100644 -index 00000000..bd1344a8 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp -@@ -0,0 +1,31 @@ -+#version 450 -+ -+#include "dequant_head.comp" -+ -+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {block_q8_0 data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; -+ -+void main() { -+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; -+ -+ const uint tid = gl_LocalInvocationID.x % 64; -+ const uint il = tid/32; -+ const uint ir = tid%32; -+ const uint ib = 32*i + ir; -+ if (ib >= p.nel / 32) { -+ return; -+ } -+ -+ const uint b_idx = 1024*i + 32*ir + 16*il; -+ -+ const float d = float(data_a[ib].d); -+ -+ const uint q_idx = 16*il; -+ -+ [[unroll]] for (uint l = 0; l < 16; l += 2) { -+ data_b[b_idx + l ] = D_TYPE(d * data_a[ib].qs[q_idx + l ]); -+ data_b[b_idx + l + 1] = D_TYPE(d * data_a[ib].qs[q_idx + l + 1]); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp -new file mode 100644 -index 00000000..4e68742b ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/diag_mask_inf.comp -@@ -0,0 +1,34 @@ -+#version 450 -+ -+#extension GL_EXT_shader_16bit_storage : require -+#extension GL_EXT_control_flow_attributes : enable -+ -+layout (push_constant) uniform parameter -+{ -+ uint ncols; -+ uint rows_per_channel; -+ uint n_past; -+} p; -+ -+#include "types.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+void main() { -+ const uint col = gl_GlobalInvocationID.y; -+ const uint row = gl_GlobalInvocationID.x; -+ -+ if (col >= p.ncols) { -+ return; -+ } -+ -+ const uint i = row*p.ncols + col; -+ if (col > p.n_past + row % p.rows_per_channel) { -+ data_d[i] = D_TYPE(uintBitsToFloat(0xFF800000)); -+ } else { -+ data_d[i] = D_TYPE(data_a[i]); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp -new file mode 100644 -index 00000000..9fb69c6c ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/div.comp -@@ -0,0 +1,27 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_binary_head.comp" -+ -+const uint num_threads = 256; -+ -+layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ uint idx = get_idx(); -+ -+ // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation -+ const uint num_iter = 2; -+ -+ [[unroll]] for (uint i = 0; i < num_iter; ++i) { -+ if (idx >= p.ne) { -+ continue; -+ } -+ uint i00, i01, i02, i03; -+ get_indices(idx, i00, i01, i02, i03); -+ -+ data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) / FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); -+ -+ idx += num_threads; -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp -new file mode 100644 -index 00000000..c5be8131 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp -@@ -0,0 +1,289 @@ -+#version 450 -+ -+#extension GL_EXT_control_flow_attributes : enable -+#extension GL_EXT_shader_16bit_storage : require -+ -+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -+#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require -+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require -+ -+#extension GL_KHR_memory_scope_semantics : enable -+#extension GL_KHR_cooperative_matrix : enable -+#extension GL_NV_cooperative_matrix2 : enable -+#extension GL_EXT_buffer_reference : enable -+#extension GL_KHR_shader_subgroup_ballot : enable -+#extension GL_KHR_shader_subgroup_vote : enable -+#extension GL_EXT_null_initializer : enable -+ -+#include "types.comp" -+#include "dequant_funcs_cm2.comp" -+ -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+layout (constant_id = 1) const uint32_t Br = 32; -+layout (constant_id = 2) const uint32_t Bc = 32; -+layout (constant_id = 3) const uint32_t D = 32; -+layout (constant_id = 4) const uint32_t Clamp = gl_CooperativeMatrixClampModeConstantNV; -+ -+layout (push_constant) uniform parameter { -+ uint32_t N; -+ uint32_t KV; -+ -+ uint32_t ne1; -+ uint32_t ne2; -+ uint32_t ne3; -+ -+ uint32_t neq2; -+ uint32_t neq3; -+ uint32_t nek2; -+ uint32_t nek3; -+ uint32_t nev2; -+ uint32_t nev3; -+ uint32_t nem1; -+ -+ uint32_t nb02; -+ uint32_t nb03; -+ uint32_t nb12; -+ uint32_t nb13; -+ uint32_t nb22; -+ uint32_t nb23; -+ uint32_t nb31; -+ -+ float scale; -+ float max_bias; -+ float logit_softcap; -+ -+ uint32_t mask; -+ uint32_t n_head_log2; -+ float m0; -+ float m1; -+} p; -+ -+layout (binding = 0) readonly buffer Q {uint8_t data_q[];}; -+layout (binding = 1) readonly buffer K {uint8_t data_k[];}; -+layout (binding = 2) readonly buffer V {uint8_t data_v[];}; -+layout (binding = 3) readonly buffer M {uint8_t data_m[];}; -+layout (binding = 4) writeonly buffer O {D_TYPE data_o[];}; -+ -+#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b)) -+ -+ACC_TYPE maxReduce(const in ACC_TYPE x, const in ACC_TYPE y) { -+ return max(x, y); -+} -+ -+ACC_TYPE smearReduce(const in ACC_TYPE x, const in ACC_TYPE y) { -+ return x; -+} -+ -+// Replace matrix elements >= numRows or numCols with 'replace' -+ACC_TYPE replacePadding(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem, const in ACC_TYPE replace, const in uint32_t numRows, const in uint32_t numCols) { -+ if (row >= numRows || col >= numCols) { -+ return replace; -+ } -+ return elem; -+} -+ -+ACC_TYPE Exp(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem) -+{ -+ return exp(elem); -+} -+ -+ACC_TYPE Max(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem0, const in ACC_TYPE elem1) -+{ -+ return max(elem0, elem1); -+} -+ -+#if defined(BLOCK_SIZE) -+#define DECODEFUNC , DEQUANTFUNC -+#else -+#define DECODEFUNC -+#endif -+ -+void main() { -+#if defined(DATA_A_IQ4_NL) -+ init_iq4nl_shmem(); -+#endif -+ -+ const uint32_t N = p.N; -+ const uint32_t KV = p.KV; -+ -+ const uint32_t Tr = CEIL_DIV(N, Br); -+ const uint32_t Tc = CEIL_DIV(KV, Bc); -+ -+ const uint32_t i = gl_WorkGroupID.x; -+ -+ const uint32_t iq2 = gl_WorkGroupID.y; -+ const uint32_t iq3 = gl_WorkGroupID.z; -+ -+ // broadcast factors -+ const uint32_t rk2 = p.neq2/p.nek2; -+ const uint32_t rk3 = p.neq3/p.nek3; -+ -+ const uint32_t rv2 = p.neq2/p.nev2; -+ const uint32_t rv3 = p.neq3/p.nev3; -+ -+ // k indices -+ const uint32_t ik3 = iq3 / rk3; -+ const uint32_t ik2 = iq2 / rk2; -+ -+ // v indices -+ const uint32_t iv3 = iq3 / rv3; -+ const uint32_t iv2 = iq2 / rv2; -+ -+ tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutQ = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); -+ tensorLayoutNV<2, Clamp> tensorLayoutK = createTensorLayoutNV(2, Clamp); -+ tensorLayoutNV<2, Clamp> tensorLayoutV = createTensorLayoutNV(2, Clamp); -+ -+ tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0); -+ -+#if defined(BLOCK_SIZE) -+ tensorLayoutK = setTensorLayoutBlockSizeNV(tensorLayoutK, 1, BLOCK_SIZE); -+ tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, BLOCK_SIZE); -+#endif -+ -+ tensorLayoutQ = setTensorLayoutDimensionNV(tensorLayoutQ, N, D); -+ tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, KV, D); -+ tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, KV, D); -+ -+ coopmat Q; -+ coopmat Qf16; -+ -+ uint32_t q_offset = iq2*p.nb02+iq3*p.nb03; -+ coopMatLoadTensorNV(Q, data_q, q_offset, sliceTensorLayoutNV(tensorLayoutQ, i * Br, Br, 0, D)); -+ -+ Qf16 = coopmat(Q); -+ Qf16 *= float16_t(p.scale); -+ -+ coopmat O = coopmat(0); -+ -+ coopmat L, M; -+ -+ L = coopmat(0); -+ M = coopmat(-1.0/0.0); -+ -+ ACC_TYPE slope = ACC_TYPE(1.0); -+ -+ // ALiBi -+ if (p.max_bias > 0.0f) { -+ const uint32_t h = iq2; -+ -+ const ACC_TYPE base = ACC_TYPE(h < p.n_head_log2 ? p.m0 : p.m1); -+ const int exph = int(h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1); -+ -+ slope = pow(base, ACC_TYPE(exph)); -+ } -+ -+ [[dont_unroll]] -+ for (uint32_t j = 0; j < Tc; ++j) { -+ -+ coopmat S = coopmat(0); -+ -+ coopmat K_T; -+ -+ uint32_t k_offset = ik2*p.nb12 + ik3*p.nb13; -+ coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, D), tensorViewTranspose DECODEFUNC); -+ S = coopMatMulAdd(Qf16, K_T, S); -+ -+ if (p.logit_softcap != 0.0f) { -+ [[unroll]] -+ for (int k = 0; k < S.length(); ++k) { -+ S[k] = ACC_TYPE(p.logit_softcap)*tanh(S[k]); -+ } -+ } -+ -+ if (p.mask != 0) { -+ tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutM = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); -+ tensorLayoutM = setTensorLayoutDimensionNV(tensorLayoutM, p.nem1, KV); -+ -+ coopmat mv; -+ -+ coopMatLoadTensorNV(mv, data_m, 0, sliceTensorLayoutNV(tensorLayoutM, i * Br, Br, j * Bc, Bc)); -+ -+ S += slope*coopmat(mv); -+ } -+ -+ // Clear padding elements to -inf, so they don't contribute to rowmax -+ if (Clamp != 0 && -+ ((j + 1) * Bc > KV || -+ (i + 1) * Br > N)) { -+ -+ uint R = ((i + 1) * Br > N) ? (N % Br) : Br; -+ uint C = ((j + 1) * Bc > KV) ? (KV % Bc) : Bc; -+ -+ coopMatPerElementNV(S, S, replacePadding, ACC_TYPE(-1.0/0.0), R, C); -+ } -+ -+ coopmat rowmax, P, rowsum, eM; -+ -+ coopMatReduceNV(rowmax, S, gl_CooperativeMatrixReduceRowNV, maxReduce); -+ -+ coopmat Mold = M; -+ -+ // M = max(rowmax, Mold) -+ // P = e^(S - M) -+ // eM = e^(Mold - M) -+ coopMatPerElementNV(M, rowmax, Max, Mold); -+ coopMatPerElementNV(P, S - M, Exp); -+ coopMatPerElementNV(eM, Mold - M, Exp); -+ -+ // Clear padding elements to 0, so they don't contribute to rowsum -+ if (Clamp != 0 && -+ ((j + 1) * Bc > KV || -+ (i + 1) * Br > N)) { -+ -+ uint R = ((i + 1) * Br > N) ? (N % Br) : Br; -+ uint C = ((j + 1) * Bc > KV) ? (KV % Bc) : Bc; -+ -+ coopMatPerElementNV(P, P, replacePadding, ACC_TYPE(0.0), R, C); -+ } -+ -+ coopmat P_A = coopmat(P); -+ -+ // compute rowsum by multiplying by matrix of all ones. -+ coopmat One = coopmat(1.0); -+ -+ rowsum = coopmat(0.0); -+ rowsum = coopMatMulAdd(P_A, One, rowsum); -+ -+ coopmat V; -+ uint32_t v_offset = iv2*p.nb22 + iv3*p.nb23; -+ coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, D) DECODEFUNC); -+ -+ L = eM*L + rowsum; -+ -+ // This is the "diagonal" matrix in the paper, but since we do componentwise -+ // multiply rather than matrix multiply it has the diagonal element smeared -+ // across the row -+ coopmat eMdiag; -+ -+ // resize eM by using smear/reduce -+ coopMatReduceNV(eMdiag, eM, gl_CooperativeMatrixReduceRowNV, smearReduce); -+ -+ O = eMdiag * O; -+ -+ O = coopMatMulAdd(P_A, V, O); -+ } -+ -+ coopmat Ldiag; -+ -+ // resize L by using smear/reduce -+ coopMatReduceNV(Ldiag, L, gl_CooperativeMatrixReduceRowNV, smearReduce); -+ -+ [[unroll]] -+ for (int k = 0; k < Ldiag.length(); ++k) { -+ Ldiag[k] = ACC_TYPE(1.0) / Ldiag[k]; -+ } -+ -+ O = Ldiag*O; -+ -+ tensorLayoutNV<3, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(3, gl_CooperativeMatrixClampModeConstantNV); -+ tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.ne2, p.ne1, D); -+ -+ // permute dimensions -+ tensorViewNV<3, false, 1, 0, 2> tensorViewPermute = createTensorViewNV(3, false, 1, 0, 2); -+ uint32_t o_offset = iq3*p.ne2*p.ne1; -+ -+ coopmat O_D = coopmat(O); -+ coopMatStoreTensorNV(O_D, data_o, o_offset, sliceTensorLayoutNV(tensorLayoutD, i * Br, Br, iq2, 1, 0, D), tensorViewPermute); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp -new file mode 100644 -index 00000000..4cc7a68c ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu.comp -@@ -0,0 +1,25 @@ -+#version 450 -+ -+#include "generic_head.comp" -+#include "types.comp" -+ -+#extension GL_EXT_control_flow_attributes : enable -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+void main() { -+ const float GELU_COEF_A = 0.044715f; -+ const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; -+ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; -+ -+ if (i >= p.KX) { -+ return; -+ } -+ -+ const float xi = float(data_a[i]); -+ const float val = SQRT_2_OVER_PI*xi*(1.0f + GELU_COEF_A*xi*xi); -+ data_d[i] = D_TYPE(0.5f*xi*(2.0f - 2.0f / (exp(2 * val) + 1))); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp -new file mode 100644 -index 00000000..e6e6fcfd ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/gelu_quick.comp -@@ -0,0 +1,23 @@ -+#version 450 -+ -+#include "generic_head.comp" -+#include "types.comp" -+ -+#extension GL_EXT_control_flow_attributes : enable -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+void main() { -+ const float GELU_QUICK_COEF = -1.702f; -+ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; -+ -+ if (i >= p.KX) { -+ return; -+ } -+ -+ const float x = float(data_a[i]); -+ data_d[i] = D_TYPE(x * (1.0f / (1.0f + exp(GELU_QUICK_COEF * x)))); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp -new file mode 100644 -index 00000000..062e2a4c ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_binary_head.comp -@@ -0,0 +1,64 @@ -+#extension GL_EXT_shader_16bit_storage : require -+#extension GL_EXT_control_flow_attributes : require -+ -+layout (push_constant) uniform parameter -+{ -+ uint ne; -+ uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03; -+ uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13; -+ uint ne20; uint ne21; uint ne22; uint ne23; uint nb20; uint nb21; uint nb22; uint nb23; -+ uint misalign_offsets; -+ float param1; float param2; int param3; -+} p; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; -+layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; -+ -+// true if src0/src1 are the same shape and the indices can be reused without additional modulus -+layout(constant_id = 0) const bool norepeat = false; -+ -+uint get_idx() { -+ return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; -+} -+ -+uint get_aoffset() { return p.misalign_offsets >> 16; } -+uint get_boffset() { return (p.misalign_offsets >> 8) & 0xFF; } -+uint get_doffset() { return p.misalign_offsets & 0xFF; } -+ -+// mod and div are expensive and coordinates/dimensions are often power of 2 or equal to 1 -+uint fastmod(uint a, uint b) { -+ if ((b & (b-1)) == 0) { -+ return a & (b-1); -+ } -+ return a % b; -+} -+ -+uint fastdiv(uint a, uint b) { -+ return (a < b) ? 0 : (a / b); -+} -+ -+void get_indices(uint idx, out uint i00, out uint i01, out uint i02, out uint i03) { -+ i03 = fastdiv(idx, (p.ne02*p.ne01*p.ne00)); -+ const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00; -+ i02 = fastdiv((idx - i03_offset), (p.ne01*p.ne00)); -+ const uint i02_offset = i02*p.ne01*p.ne00; -+ i01 = (idx - i03_offset - i02_offset) / p.ne00; -+ i00 = idx - i03_offset - i02_offset - i01*p.ne00; -+} -+ -+uint src0_idx(uint i00, uint i01, uint i02, uint i03) { -+ return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00; -+} -+ -+uint src1_idx(uint i00, uint i01, uint i02, uint i03) { -+ if (norepeat) { -+ return i03*p.nb13 + i02*p.nb12 + i01*p.nb11 + i00*p.nb10; -+ } else { -+ return fastmod(i03, p.ne13)*p.nb13 + fastmod(i02, p.ne12)*p.nb12 + fastmod(i01, p.ne11)*p.nb11 + fastmod(i00, p.ne10)*p.nb10; -+ } -+} -+ -+uint dst_idx(uint i00, uint i01, uint i02, uint i03) { -+ return i03*p.nb23 + i02*p.nb22 + i01*p.nb21 + i00*p.nb20; -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp -new file mode 100644 -index 00000000..66e46ae6 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_head.comp -@@ -0,0 +1,9 @@ -+#extension GL_EXT_shader_16bit_storage : require -+ -+layout (push_constant) uniform parameter -+{ -+ uint KX; -+ uint KY; -+ float param1; -+ float param2; -+} p; -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp -new file mode 100644 -index 00000000..68d1bc9f ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.comp -@@ -0,0 +1,56 @@ -+#extension GL_EXT_shader_16bit_storage : require -+#extension GL_EXT_control_flow_attributes : require -+ -+layout (push_constant) uniform parameter -+{ -+ uint ne; -+ uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03; -+ uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13; -+ uint misalign_offsets; -+ float param1; float param2; -+ -+ uint ne0_012mp; uint ne0_012L; -+ uint ne0_01mp; uint ne0_01L; -+ uint ne0_0mp; uint ne0_0L; -+ uint ne1_012mp; uint ne1_012L; -+ uint ne1_01mp; uint ne1_01L; -+ uint ne1_0mp; uint ne1_0L; -+} p; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+uint get_idx() { -+ return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; -+} -+ -+uint get_aoffset() { return p.misalign_offsets >> 16; } -+uint get_doffset() { return p.misalign_offsets & 0xFFFF; } -+ -+// see init_fastdiv_values in ggml-vulkan.cpp -+uint fastdiv(uint n, uint mp, uint L) { -+ uint msbs, lsbs; -+ // msbs = mulhi(n, mp) -+ umulExtended(n, mp, msbs, lsbs); -+ return (msbs + n) >> L; -+} -+ -+uint src0_idx(uint idx) { -+ const uint i03 = fastdiv(idx, p.ne0_012mp, p.ne0_012L); -+ const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00; -+ const uint i02 = fastdiv(idx - i03_offset, p.ne0_01mp, p.ne0_01L); -+ const uint i02_offset = i02*p.ne01*p.ne00; -+ const uint i01 = fastdiv(idx - i03_offset - i02_offset, p.ne0_0mp, p.ne0_0L); -+ const uint i00 = idx - i03_offset - i02_offset - i01*p.ne00; -+ return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00; -+} -+ -+uint dst_idx(uint idx) { -+ const uint i13 = fastdiv(idx, p.ne1_012mp, p.ne1_012L); -+ const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10; -+ const uint i12 = fastdiv(idx - i13_offset, p.ne1_01mp, p.ne1_01L); -+ const uint i12_offset = i12*p.ne11*p.ne10; -+ const uint i11 = fastdiv(idx - i13_offset - i12_offset, p.ne1_0mp, p.ne1_0L); -+ const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10; -+ return i13*p.nb13 + i12*p.nb12 + i11*p.nb11 + i10*p.nb10; -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp -new file mode 100644 -index 00000000..e877ed77 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows.comp -@@ -0,0 +1,28 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_binary_head.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ const uint i00 = gl_GlobalInvocationID.x; -+ const uint i10 = gl_GlobalInvocationID.y; -+ const uint i11 = (gl_GlobalInvocationID.z)/p.ne12; -+ const uint i12 = (gl_GlobalInvocationID.z)%p.ne12; -+ -+ if (i00 >= p.ne00) { -+ return; -+ } -+ -+ const uint i01 = data_b[get_boffset() + i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; -+ -+ const uint a_offset = get_aoffset() + i01*p.nb01 + i11*p.nb02 + i12*p.nb03; -+ const uint d_offset = get_doffset() + i10*p.nb21 + i11*p.nb22 + i12*p.nb23; -+ -+#ifndef OPTIMIZATION_ERROR_WORKAROUND -+ data_d[d_offset + i00] = D_TYPE(data_a[a_offset + i00]); -+#else -+ data_d[d_offset + i00] = data_a[a_offset + i00]; -+#endif -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp -new file mode 100644 -index 00000000..1426fde6 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp -@@ -0,0 +1,39 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_binary_head.comp" -+#include "dequant_funcs.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ const uint i00 = (gl_GlobalInvocationID.x)*2; -+ const uint i10 = gl_GlobalInvocationID.y; -+ const uint i11 = (gl_GlobalInvocationID.z)/p.ne12; -+ const uint i12 = (gl_GlobalInvocationID.z)%p.ne12; -+ -+#if defined(DATA_A_IQ4_NL) -+ init_iq4nl_shmem(); -+#endif -+ -+ if (i00 >= p.ne00) { -+ return; -+ } -+ -+ const uint i01 = data_b[i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; -+ -+ const uint a_offset = i01*p.nb01 + i11*p.nb02 + i12*p.nb03; -+ const uint d_offset = i10*p.nb21 + i11*p.nb22 + i12*p.nb23; -+ -+ const uint ib = a_offset + i00/QUANT_K; // block index -+ const uint iqs = (i00%QUANT_K)/QUANT_R; // quant index -+ const uint iybs = i00 - i00%QUANT_K; // dst block start index -+ const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2; -+ -+ vec2 v = dequantize(ib, iqs, 0); -+ const vec2 dm = get_dm(ib, 0); -+ v = v * dm.x + dm.y; -+ -+ data_d[d_offset + iybs + iqs ] = D_TYPE(v.x); -+ data_d[d_offset + iybs + iqs + y_offset] = D_TYPE(v.y); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp -new file mode 100644 -index 00000000..b6a0d564 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/group_norm.comp -@@ -0,0 +1,66 @@ -+#version 450 -+ -+#include "generic_head.comp" -+#include "types.comp" -+ -+#extension GL_EXT_control_flow_attributes : enable -+#define BLOCK_SIZE 512 -+ -+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+shared float tmp[BLOCK_SIZE]; -+ -+void main() { -+ const uint group_size = p.KX; -+ const float eps = p.param1; -+ -+ const uint tid = gl_LocalInvocationID.x; -+ const uint start = gl_WorkGroupID.x * group_size + tid; -+ const uint end = (gl_WorkGroupID.x + 1) * group_size; -+ -+ tmp[tid] = 0.0f; -+ -+ // Calculate mean -+ [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { -+ tmp[tid] += float(data_a[col]); -+ } -+ -+ // tmp up partial tmps and write back result -+ barrier(); -+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { -+ if (tid < s) { -+ tmp[tid] += tmp[tid + s]; -+ } -+ barrier(); -+ } -+ -+ const float mean = tmp[0] / group_size; -+ barrier(); -+ tmp[tid] = 0.0f; -+ -+ // Calculate variance -+ [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { -+ const float xi = float(data_a[col]) - mean; -+ data_d[col] = D_TYPE(xi); -+ tmp[tid] += xi * xi; -+ } -+ -+ // sum up partial sums and write back result -+ barrier(); -+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { -+ if (tid < s) { -+ tmp[tid] += tmp[tid + s]; -+ } -+ barrier(); -+ } -+ -+ const float variance = tmp[0] / group_size; -+ const float scale = inversesqrt(variance + eps); -+ -+ [[unroll]] for (uint col = start; col < end; col += BLOCK_SIZE) { -+ data_d[col] *= D_TYPE(scale); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp -new file mode 100644 -index 00000000..122b1e93 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp -@@ -0,0 +1,87 @@ -+#version 450 -+ -+#extension GL_EXT_shader_16bit_storage : require -+#extension GL_EXT_spirv_intrinsics: enable -+#extension GL_EXT_control_flow_attributes : require -+ -+#if RTE16 -+spirv_execution_mode(capabilities = [4467], 4462, 16); // RoundingModeRTE, 16 bits -+#endif -+ -+layout (push_constant) uniform parameter -+{ -+ uint batch_offset; uint offset_delta; -+ uint IC; -+ uint IW; uint IH; -+ uint OW; uint OH; -+ uint KW; uint KH; -+ uint pelements; -+ uint CHW; -+ int s0; int s1; -+ int p0; int p1; -+ int d0; int d1; -+} p; -+ -+#include "types.comp" -+ -+layout(constant_id = 0) const uint BLOCK_SIZE = 32; -+ -+const uint NUM_ITER = 512 / BLOCK_SIZE; -+ -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+void main() { -+ const uint gidx = gl_GlobalInvocationID.x; -+ -+ const uint oh = gl_GlobalInvocationID.y; -+ const uint batch = gl_GlobalInvocationID.z / p.IC; -+ const uint ic = gl_GlobalInvocationID.z % p.IC; -+ -+ A_TYPE values[NUM_ITER]; -+ uint offset_dst[NUM_ITER]; -+ [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { -+ values[idx] = A_TYPE(0); -+ } -+ -+ [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { -+ -+ const uint i = gidx * NUM_ITER + idx; -+ -+ const uint ksize = p.OW * (p.KH > 1 ? p.KW : 1); -+ const uint kx = i / ksize; -+ const uint kd = kx * ksize; -+ const uint ky = (i - kd) / p.OW; -+ const uint ix = i % p.OW; -+ -+ const uint iiw = ix * p.s0 + kx * p.d0 - p.p0; -+ const uint iih = oh * p.s1 + ky * p.d1 - p.p1; -+ -+ offset_dst[idx] = -+ ((batch * p.OH + oh) * p.OW + ix) * p.CHW + -+ (ic * (p.KW * p.KH) + ky * p.KW + kx); -+ -+ if (i >= p.pelements) { -+ continue; -+ } -+ -+ if (iih < p.IH && iiw < p.IW) { -+ const uint offset_src = ic * p.offset_delta + batch * p.batch_offset; -+ values[idx] = data_a[offset_src + iih * p.IW + iiw]; -+ } -+ } -+ -+ [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) { -+ -+ const uint i = gidx * NUM_ITER + idx; -+ -+ if (i >= p.pelements) { -+ continue; -+ } -+ -+ data_d[offset_dst[idx]] = D_TYPE(values[idx]); -+ } -+ -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp -new file mode 100644 -index 00000000..d90a99ae ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/leaky_relu.comp -@@ -0,0 +1,22 @@ -+#version 450 -+ -+#include "generic_head.comp" -+#include "types.comp" -+ -+#extension GL_EXT_control_flow_attributes : enable -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+void main() { -+ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; -+ -+ if (i >= p.KX) { -+ return; -+ } -+ -+ const float val = float(data_a[i]); -+ data_d[i] = D_TYPE(max(val, 0.0f) + min(val, 0.0f) * p.param1); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp -new file mode 100644 -index 00000000..43de19df ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul.comp -@@ -0,0 +1,27 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_binary_head.comp" -+ -+const uint num_threads = 256; -+ -+layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ uint idx = get_idx(); -+ -+ // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation -+ const uint num_iter = 2; -+ -+ [[unroll]] for (uint i = 0; i < num_iter; ++i) { -+ if (idx >= p.ne) { -+ continue; -+ } -+ uint i00, i01, i02, i03; -+ get_indices(idx, i00, i01, i02, i03); -+ -+ data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) * FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); -+ -+ idx += num_threads; -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp -new file mode 100644 -index 00000000..4c64fd47 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_split_k_reduce.comp -@@ -0,0 +1,48 @@ -+#version 450 -+ -+#extension GL_EXT_control_flow_attributes : enable -+ -+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {float data_a[];}; -+layout (binding = 0) readonly buffer A4 {vec4 data_a4[];}; -+layout (binding = 1) writeonly buffer D {float data_d[];}; -+layout (binding = 1) writeonly buffer D4 {vec4 data_d4[];}; -+ -+layout (push_constant) uniform parameter { -+ uint ne; -+ uint k_num; -+} p; -+ -+void main() { -+ // Each invocation handles four consecutive components -+ const uint idx = gl_GlobalInvocationID.x * 4; -+ -+ if (idx >= p.ne) { -+ return; -+ } -+ -+ // Check if all four components are in bounds and aligned, -+ // then use vector loads -+ if (idx + 3 < p.ne && (p.ne % 4) == 0) { -+ vec4 result = vec4(0.0f); -+ -+ [[unroll]] for (uint i = 0; i < p.k_num; i++) { -+ result += data_a4[(i * p.ne + idx) / 4]; -+ } -+ -+ data_d4[idx / 4] = result; -+ } else { -+ [[unroll]] for (uint j = 0; j < 4; ++j) { -+ if (idx + j < p.ne) { -+ float result = 0.0f; -+ -+ [[unroll]] for (uint i = 0; i < p.k_num; i++) { -+ result += data_a[i * p.ne + idx + j]; -+ } -+ -+ data_d[idx + j] = result; -+ } -+ } -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp -new file mode 100644 -index 00000000..24875cdc ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp -@@ -0,0 +1,152 @@ -+#version 450 -+ -+#ifdef FLOAT16 -+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -+#endif -+#extension GL_EXT_shader_explicit_arithmetic_types : require -+ -+#include "mul_mat_vec_base.comp" -+ -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+#if !defined(DATA_A_F32) && !defined(DATA_A_F16) -+#define K_PER_ITER 8 -+#else -+#define K_PER_ITER 2 -+#endif -+ -+ -+uint a_offset, b_offset, d_offset, y_offset; -+ -+void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint tid, const uint i, bool lastiter) -+{ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ const uint col = i*BLOCK_SIZE + K_PER_ITER*tid; -+ const uint iqs = (col%QUANT_K)/QUANT_R; // quant index -+ const uint iybs = col - col%QUANT_K; // y block start index -+ -+#if K_PER_ITER == 8 -+#if QUANT_R == 2 -+ const B_TYPE_VEC4 bv02 = data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]; -+ const B_TYPE_VEC4 bv13 = data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs + y_offset) / 4]; -+ const vec4 bv0 = vec4(bv02.x, bv13.x, bv02.y, bv13.y); -+ const vec4 bv1 = vec4(bv02.z, bv13.z, bv02.w, bv13.w); -+#else -+ const vec4 bv0 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]); -+ const vec4 bv1 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4 + 1]); -+#endif -+#else -+ // Check if the second of the pair of elements is OOB, and don't fetch B or -+ // accumulate it. We still fetch a pair of elements for A, which is fine for -+ // quantized formats since they'll be within the same block. We should -+ // probably skip fetching the second element for F16/F32, but as of now we -+ // still do. -+ const bool OOB = lastiter && (iybs + iqs + y_offset >= p.ncols); -+ -+ FLOAT_TYPE b0 = 0, b1 = 0; -+ b0 = FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs]); -+ if (!OOB) { -+ b1 = FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs + y_offset]); -+ } -+#endif -+ uint ibi = first_row*p.ncols; -+ [[unroll]] for (uint n = 0; n < num_rows; ++n) { -+ const uint ib = (ibi + col)/QUANT_K; // block index -+ ibi += p.ncols; -+ -+#if K_PER_ITER == 8 -+ vec4 v = dequantize4(ib, iqs, a_offset); -+ vec4 v2 = dequantize4(ib, iqs+(4/QUANT_R), a_offset); -+ -+ const vec2 dm = get_dm(ib, a_offset); -+ if (dm.y != 0) { // quant has min component -+ v = v * dm.x + dm.y; -+ v2 = v2 * dm.x + dm.y; -+ } -+ -+ // matrix multiplication -+ FLOAT_TYPE rowtmp = dot(bv0, v); -+ rowtmp += dot(bv1, v2); -+ -+ if (dm.y == 0) -+ rowtmp *= dm.x; -+ -+ temp[j][n] += rowtmp; -+#else -+ const vec2 v = dequantize(ib, iqs, a_offset); -+ -+ // matrix multiplication -+ temp[j][n] = fma(FLOAT_TYPE(v.x), b0, temp[j][n]); -+ if (!OOB) { -+ temp[j][n] = fma(FLOAT_TYPE(v.y), b1, temp[j][n]); -+ } -+#endif -+ } -+ } -+} -+ -+void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { -+ const uint tid = gl_LocalInvocationID.x; -+ -+ get_offsets(a_offset, b_offset, d_offset); -+ a_offset /= QUANT_K; -+ -+ y_offset = QUANT_R == 1 ? 1 : QUANT_K/2; -+ -+ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; -+ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { -+ temp[j][i] = FLOAT_TYPE(0); -+ } -+ } -+ -+ uint num_iters = p.ncols / (K_PER_ITER * BLOCK_SIZE); -+ if (num_iters * K_PER_ITER * BLOCK_SIZE + K_PER_ITER*tid < p.ncols) { -+ num_iters++; -+ } -+ int unroll_count = 4; -+ uint unrolled_iters = num_iters & ~(unroll_count - 1); -+ -+ uint i = 0; -+ while (i < unrolled_iters) { -+ // Manually partially unroll the loop -+ [[unroll]] for (uint k = 0; k < unroll_count; ++k) { -+ iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false); -+ i++; -+ } -+ } -+ unroll_count = 2; -+ unrolled_iters = num_iters & ~(unroll_count - 1); -+ while (i < unrolled_iters) { -+ // Manually partially unroll the loop -+ [[unroll]] for (uint k = 0; k < unroll_count; ++k) { -+ iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false); -+ i++; -+ } -+ } -+ while (i < num_iters) { -+ iter(temp, first_row, num_rows, tid, i*K_PER_ITER, true); -+ i++; -+ } -+ -+ reduce_result(temp, d_offset, first_row, num_rows, tid); -+} -+ -+void main() { -+ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); -+ -+#if defined(DATA_A_IQ4_NL) -+ init_iq4nl_shmem(); -+#endif -+ -+ // do NUM_ROWS at a time, unless there aren't enough remaining rows -+ if (first_row + NUM_ROWS <= p.stride_d) { -+ compute_outputs(first_row, NUM_ROWS); -+ } else { -+ if (first_row >= p.stride_d) { -+ return; -+ } -+ compute_outputs(first_row, p.stride_d - first_row); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp -new file mode 100644 -index 00000000..903753c7 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.comp -@@ -0,0 +1,118 @@ -+#extension GL_EXT_control_flow_attributes : enable -+#extension GL_EXT_shader_16bit_storage : require -+#extension GL_EXT_shader_8bit_storage : require -+ -+#ifdef MUL_MAT_ID -+#define EXPERT_COUNT 8 -+#endif -+ -+#include "types.comp" -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; -+layout (binding = 1) readonly buffer BV2 {B_TYPE_VEC2 data_b_v2[];}; -+layout (binding = 1) readonly buffer BV4 {B_TYPE_VEC4 data_b_v4[];}; -+ -+layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; -+#ifdef MUL_MAT_ID -+layout (binding = 3) readonly buffer IDS {int data_ids[];}; -+#endif -+ -+#include "dequant_funcs.comp" -+ -+layout (push_constant) uniform parameter -+{ -+ uint ncols; -+ uint stride_a; -+ uint stride_b; -+ uint stride_d; -+ -+ uint batch_stride_a; -+ uint batch_stride_b; -+ uint batch_stride_d; -+ -+#ifdef MUL_MAT_ID -+ uint nei0; -+ uint ne11; -+#else -+ uint ne02; -+ uint ne12; -+ uint broadcast2; -+ uint broadcast3; -+#endif -+} p; -+ -+void get_offsets(out uint a_offset, out uint b_offset, out uint d_offset) { -+#ifdef MUL_MAT_ID -+ const uint expert_idx = gl_GlobalInvocationID.y; -+#else -+ const uint batch_idx = gl_GlobalInvocationID.y; -+#endif -+ -+#ifndef MUL_MAT_ID -+ uint batch_idx_a = 0; -+ if (batch_idx != 0) { -+ const uint i13 = batch_idx / p.ne12; -+ const uint i12 = batch_idx % p.ne12; -+ -+ const uint i03 = i13 / p.broadcast3; -+ const uint i02 = i12 / p.broadcast2; -+ -+ batch_idx_a = i03 * p.ne02 + i02; -+ } -+#else -+ const uint expert_id = data_ids[expert_idx]; -+#endif -+ -+ a_offset = -+#ifdef MUL_MAT_ID -+ expert_id * p.batch_stride_a; -+#else -+ batch_idx_a * p.batch_stride_a; -+#endif -+ b_offset = -+#ifdef MUL_MAT_ID -+ (expert_idx % p.ne11) * p.stride_b; -+#else -+ batch_idx * p.batch_stride_b; -+#endif -+ d_offset = -+#ifdef MUL_MAT_ID -+ expert_idx * p.stride_d; -+#else -+ batch_idx * p.batch_stride_d; -+#endif -+} -+ -+layout (constant_id = 0) const uint BLOCK_SIZE = 32; -+layout (constant_id = 1) const uint NUM_ROWS = 1; -+layout (constant_id = 2) const uint NUM_COLS = 1; -+ -+shared FLOAT_TYPE tmpsh[NUM_COLS][NUM_ROWS][BLOCK_SIZE]; -+ -+void reduce_result(const in FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const in uint32_t d_offset, const in uint32_t first_row, const in uint32_t num_rows, const in uint32_t tid) { -+ // sum up partial sums and write back result -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ [[unroll]] for (uint n = 0; n < num_rows; ++n) { -+ tmpsh[j][n][tid] = temp[j][n]; -+ } -+ } -+ barrier(); -+ [[unroll]] for (uint s = BLOCK_SIZE/2; s > 0; s >>= 1) { -+ if (tid < s) { -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ [[unroll]] for (uint n = 0; n < num_rows; ++n) { -+ tmpsh[j][n][tid] += tmpsh[j][n][tid + s]; -+ } -+ } -+ } -+ barrier(); -+ } -+ if (tid == 0) { -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ [[unroll]] for (uint n = 0; n < num_rows; ++n) { -+ data_d[j*p.batch_stride_d + d_offset + first_row + n] = D_TYPE(tmpsh[j][n][0]); -+ } -+ } -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp -new file mode 100644 -index 00000000..1cc4996d ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp -@@ -0,0 +1,71 @@ -+#version 450 -+ -+#extension GL_EXT_control_flow_attributes : enable -+#extension GL_EXT_shader_16bit_storage : require -+ -+#define BLOCK_SIZE 32 -+#define FLOAT_TYPE float -+ -+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; -+layout (binding = 2) writeonly buffer D {D_TYPE dst[];}; -+ -+layout (push_constant) uniform parameter -+{ -+ uint ncols_x; -+ uint nrows_x; -+ uint row_stride_x; -+ uint channel_stride_x; -+ uint channel_x_divisor; -+ uint b_offset; -+ uint d_offset; -+} p; -+ -+shared FLOAT_TYPE tmp[BLOCK_SIZE]; -+ -+void main() { -+ const uint tid = gl_LocalInvocationID.x; -+ const uint row_x = gl_GlobalInvocationID.y; -+ const uint channel = gl_GlobalInvocationID.z; -+ const uint channel_x = channel / p.channel_x_divisor; -+ -+ const uint nrows_y = p.ncols_x; -+ const uint nrows_dst = p.nrows_x; -+ const uint row_dst = row_x; -+ -+ const uint idst = channel*nrows_dst + row_dst; -+ -+ tmp[tid] = 0.0f; -+ -+ for (uint col_x0 = 0; col_x0 < p.ncols_x; col_x0 += BLOCK_SIZE) { -+ const uint col_x = col_x0 + tid; -+ -+ if (col_x >= p.ncols_x) { -+ break; -+ } -+ -+ const uint row_y = col_x; -+ -+ const uint ix = channel_x*p.channel_stride_x + row_x*p.row_stride_x + col_x; -+ const uint iy = channel*nrows_y + row_y; -+ -+ const FLOAT_TYPE xi = FLOAT_TYPE(data_a[ix]); -+ -+ tmp[tid] = fma(xi, FLOAT_TYPE(data_b[iy]), tmp[tid]); -+ } -+ -+ // sum up partial sums and write back result -+ barrier(); -+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { -+ if (tid < s) { -+ tmp[tid] += tmp[tid + s]; -+ } -+ barrier(); -+ } -+ -+ if (tid == 0) { -+ dst[idst] = tmp[0]; -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp -new file mode 100644 -index 00000000..9b443807 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_p021.comp -@@ -0,0 +1,73 @@ -+#version 450 -+ -+#extension GL_EXT_control_flow_attributes : enable -+#extension GL_EXT_shader_16bit_storage : require -+ -+#define BLOCK_SIZE 32 -+#define FLOAT_TYPE float -+ -+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; -+layout (binding = 2) writeonly buffer D {D_TYPE dst[];}; -+ -+layout (push_constant) uniform parameter -+{ -+ uint ncols_x; -+ uint nrows_x; -+ uint nchannels_x; -+ uint nchannels_y; -+ uint b_offset; -+ uint d_offset; -+} p; -+ -+shared FLOAT_TYPE tmp[BLOCK_SIZE]; -+ -+void main() { -+ const uint tid = gl_LocalInvocationID.x; -+ const uint row_x = gl_GlobalInvocationID.y; -+ const uint channel = gl_GlobalInvocationID.z; -+ const uint channel_x = channel / (p.nchannels_y / p.nchannels_x); -+ -+ const uint nrows_y = p.ncols_x; -+ const uint nrows_dst = p.nrows_x; -+ const uint row_dst = row_x; -+ -+ tmp[tid] = FLOAT_TYPE(0.0f); -+ -+ for (uint col_x0 = 0; col_x0 < p.ncols_x; col_x0 += BLOCK_SIZE) { -+ const uint col_x = col_x0 + tid; -+ -+ if (col_x >= p.ncols_x) { -+ break; -+ } -+ -+ // x is transposed and permuted -+ const uint ix = row_x*p.nchannels_x*p.ncols_x + channel_x*p.ncols_x + col_x; -+ const FLOAT_TYPE xi = FLOAT_TYPE(data_a[ix]); -+ -+ const uint row_y = col_x; -+ -+ // y is not transposed but permuted -+ const uint iy = channel*nrows_y + row_y; -+ -+ tmp[tid] = fma(xi, FLOAT_TYPE(data_b[iy]), tmp[tid]); -+ } -+ -+ // dst is not transposed and not permuted -+ const uint idst = channel*nrows_dst + row_dst; -+ -+ // sum up partial sums and write back result -+ barrier(); -+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { -+ if (tid < s) { -+ tmp[tid] += tmp[tid + s]; -+ } -+ barrier(); -+ } -+ -+ if (tid == 0) { -+ dst[idst] = tmp[0]; -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp -new file mode 100644 -index 00000000..93421344 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q2_k.comp -@@ -0,0 +1,115 @@ -+#version 450 -+#extension GL_EXT_shader_explicit_arithmetic_types : require -+ -+#include "mul_mat_vec_base.comp" -+ -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { -+ uint a_offset, b_offset, d_offset; -+ get_offsets(a_offset, b_offset, d_offset); -+ -+ const uint num_blocks_per_row = p.ncols / QUANT_K; -+ -+ // 16 threads are used to process each block -+ const uint it_size = gl_WorkGroupSize.x/16; -+ const uint tid = gl_LocalInvocationID.x; -+ const uint itid = tid%16; // 0...16 -+ const uint ix = tid/16; -+ -+ const uint step = 8; -+ -+ const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... -+ const uint v_in = itid - step*v_im; // 0...15 or 0...7 -+ -+ const uint l0 = 2*v_in; // 0...15 -+ const uint q_offset = 32*v_im + l0; -+ const uint s_offset = 8*v_im; -+ const uint y_offset = 128*v_im + l0; -+ -+ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; -+ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { -+ temp[j][i] = FLOAT_TYPE(0); -+ } -+ } -+ -+ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { -+ const uint y_idx = i * QUANT_K + y_offset; -+ -+ [[unroll]] for (uint n = 0; n < num_rows; ++n) { -+ const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; -+ f16vec2 d = data_a[ib0 + i].d; -+ const FLOAT_TYPE dall = d.x; -+ const FLOAT_TYPE dmin = d.y; -+ -+ uint32_t s0_u32 = data_a_packed32[ib0 + i].scales[s_offset / 4 + 0]; -+ uint32_t s4_u32 = data_a_packed32[ib0 + i].scales[s_offset / 4 + 1]; -+ -+ uint32_t s0_lo4_u32 = s0_u32 & 0x0F0F0F0F; -+ uint32_t s0_hi4_u32 = (s0_u32 >> 4) & 0x0F0F0F0F; -+ uint32_t s4_lo4_u32 = s4_u32 & 0x0F0F0F0F; -+ uint32_t s4_hi4_u32 = (s4_u32 >> 4) & 0x0F0F0F0F; -+ -+ uvec4 s0_lo4 = uvec4(unpack8(s0_lo4_u32)); -+ uvec4 s4_lo4 = uvec4(unpack8(s4_lo4_u32)); -+ uvec4 s0_hi4 = uvec4(unpack8(s0_hi4_u32)); -+ uvec4 s4_hi4 = uvec4(unpack8(s4_hi4_u32)); -+ -+ uint16_t qs0_u16 = data_a_packed16[ib0 + i].qs[q_offset / 2 + 0]; -+ uint16_t qs16_u16 = data_a_packed16[ib0 + i].qs[q_offset / 2 + 8]; -+ uvec2 qs0 = uvec2(unpack8(qs0_u16)); -+ uvec2 qs16 = uvec2(unpack8(qs16_u16)); -+ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ B_TYPE_VEC2 b0 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 0]; -+ B_TYPE_VEC2 b16 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 8]; -+ B_TYPE_VEC2 b32 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 16]; -+ B_TYPE_VEC2 b48 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 24]; -+ B_TYPE_VEC2 b64 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 32]; -+ B_TYPE_VEC2 b80 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 40]; -+ B_TYPE_VEC2 b96 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 48]; -+ B_TYPE_VEC2 b112 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 56]; -+ -+ FLOAT_TYPE sum1 = FLOAT_TYPE(0.0); -+ FLOAT_TYPE sum2 = FLOAT_TYPE(0.0); -+ [[unroll]] for (int l = 0; l < 2; ++l) { -+ sum1 = fma(FLOAT_TYPE(b0[l]), FLOAT_TYPE(s0_lo4[0]) * FLOAT_TYPE((qs0[l] >> 0) & 3), -+ fma(FLOAT_TYPE(b16[l]), FLOAT_TYPE(s0_lo4[1]) * FLOAT_TYPE((qs16[l] >> 0) & 3), -+ fma(FLOAT_TYPE(b32[l]), FLOAT_TYPE(s0_lo4[2]) * FLOAT_TYPE((qs0[l] >> 2) & 3), -+ fma(FLOAT_TYPE(b48[l]), FLOAT_TYPE(s0_lo4[3]) * FLOAT_TYPE((qs16[l] >> 2) & 3), -+ fma(FLOAT_TYPE(b64[l]), FLOAT_TYPE(s4_lo4[0]) * FLOAT_TYPE((qs0[l] >> 4) & 3), -+ fma(FLOAT_TYPE(b80[l]), FLOAT_TYPE(s4_lo4[1]) * FLOAT_TYPE((qs16[l] >> 4) & 3), -+ fma(FLOAT_TYPE(b96[l]), FLOAT_TYPE(s4_lo4[2]) * FLOAT_TYPE((qs0[l] >> 6) & 3), -+ fma(FLOAT_TYPE(b112[l]), FLOAT_TYPE(s4_lo4[3]) * FLOAT_TYPE((qs16[l] >> 6) & 3), sum1)))))))); -+ sum2 = fma(FLOAT_TYPE(b0[l]), FLOAT_TYPE(s0_hi4[0]), -+ fma(FLOAT_TYPE(b16[l]), FLOAT_TYPE(s0_hi4[1]), -+ fma(FLOAT_TYPE(b32[l]), FLOAT_TYPE(s0_hi4[2]), -+ fma(FLOAT_TYPE(b48[l]), FLOAT_TYPE(s0_hi4[3]), -+ fma(FLOAT_TYPE(b64[l]), FLOAT_TYPE(s4_hi4[0]), -+ fma(FLOAT_TYPE(b80[l]), FLOAT_TYPE(s4_hi4[1]), -+ fma(FLOAT_TYPE(b96[l]), FLOAT_TYPE(s4_hi4[2]), -+ fma(FLOAT_TYPE(b112[l]), FLOAT_TYPE(s4_hi4[3]), sum2)))))))); -+ } -+ temp[j][n] = fma(dall, sum1, fma(-dmin, sum2, temp[j][n])); -+ } -+ } -+ } -+ -+ reduce_result(temp, d_offset, first_row, num_rows, tid); -+} -+ -+void main() { -+ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); -+ -+ // do NUM_ROWS at a time, unless there aren't enough remaining rows -+ if (first_row + NUM_ROWS <= p.stride_d) { -+ compute_outputs(first_row, NUM_ROWS); -+ } else { -+ if (first_row >= p.stride_d) { -+ return; -+ } -+ compute_outputs(first_row, p.stride_d - first_row); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp -new file mode 100644 -index 00000000..86b0159d ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp -@@ -0,0 +1,103 @@ -+#version 450 -+#extension GL_EXT_shader_explicit_arithmetic_types : require -+ -+#include "mul_mat_vec_base.comp" -+ -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { -+ uint a_offset, b_offset, d_offset; -+ get_offsets(a_offset, b_offset, d_offset); -+ -+ const uint num_blocks_per_row = p.ncols / QUANT_K; -+ -+ // 16 threads are used to process each block -+ const uint it_size = gl_WorkGroupSize.x/16; -+ const uint tid = gl_LocalInvocationID.x; -+ const uint itid = tid%16; // 0...16 -+ const uint ix = tid/16; -+ -+ const uint step = 8; -+ -+ const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... -+ const uint v_in = itid - step*v_im; // 0...15 or 0...7 -+ -+ const uint8_t m = uint8_t(1 << (4 * v_im)); -+ -+ const uint l0 = 2*v_in; // 0...15 -+ const uint q_offset = 32*v_im + l0; -+ const uint y_offset = 128*v_im + l0; -+ -+ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; -+ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { -+ temp[j][i] = FLOAT_TYPE(0); -+ } -+ } -+ -+ const uint s_shift = 4 * v_im; -+ -+ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { -+ const uint y_idx = i * QUANT_K + y_offset; -+ -+ [[unroll]] for (uint n = 0; n < num_rows; ++n) { -+ const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; -+ const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d); -+ -+ uint16_t s0_16 = data_a_packed16[ib0 + i].scales[0]; -+ uint16_t s2_16 = data_a_packed16[ib0 + i].scales[1]; -+ uint16_t s4_16 = data_a_packed16[ib0 + i].scales[2]; -+ uint16_t s6_16 = data_a_packed16[ib0 + i].scales[3]; -+ uint16_t s8_16 = data_a_packed16[ib0 + i].scales[4]; -+ uint16_t s10_16 = data_a_packed16[ib0 + i].scales[5]; -+ u8vec2 s0 = unpack8(s0_16); -+ u8vec2 s2 = unpack8(s2_16); -+ u8vec2 s4 = unpack8(s4_16); -+ u8vec2 s6 = unpack8(s6_16); -+ u8vec2 s8 = unpack8(s8_16); -+ u8vec2 s10 = unpack8(s10_16); -+ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ -+ B_TYPE_VEC2 b0 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 0]; -+ B_TYPE_VEC2 b16 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 8]; -+ B_TYPE_VEC2 b32 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 16]; -+ B_TYPE_VEC2 b48 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 24]; -+ B_TYPE_VEC2 b64 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 32]; -+ B_TYPE_VEC2 b80 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 40]; -+ B_TYPE_VEC2 b96 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 48]; -+ B_TYPE_VEC2 b112 = data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 56]; -+ -+ FLOAT_TYPE sum = FLOAT_TYPE(0.0); -+ [[unroll]] for (int l = 0; l < 2; ++l) { -+ sum = fma(FLOAT_TYPE(b0[l]) * FLOAT_TYPE(int8_t(((s0[0] >> s_shift) & 0xF) | ((s8[0] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 0)) != 0) ? 0 : 4)), -+ fma(FLOAT_TYPE(b32[l]) * FLOAT_TYPE(int8_t(((s2[0] >> s_shift) & 0xF) | ((s10[0] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 1)) != 0) ? 0 : 4)), -+ fma(FLOAT_TYPE(b64[l]) * FLOAT_TYPE(int8_t(((s4[0] >> s_shift) & 0xF) | ((s8[0] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 2)) != 0) ? 0 : 4)), -+ fma(FLOAT_TYPE(b96[l]) * FLOAT_TYPE(int8_t(((s6[0] >> s_shift) & 0xF) | ((s10[0] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 3)) != 0) ? 0 : 4)), -+ fma(FLOAT_TYPE(b16[l]) * FLOAT_TYPE(int8_t(((s0[1] >> s_shift) & 0xF) | ((s8[1] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 0)) != 0) ? 0 : 4)), -+ fma(FLOAT_TYPE(b48[l]) * FLOAT_TYPE(int8_t(((s2[1] >> s_shift) & 0xF) | ((s10[1] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 1)) != 0) ? 0 : 4)), -+ fma(FLOAT_TYPE(b80[l]) * FLOAT_TYPE(int8_t(((s4[1] >> s_shift) & 0xF) | ((s8[1] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 2)) != 0) ? 0 : 4)), -+ fma(FLOAT_TYPE(b112[l]) * FLOAT_TYPE(int8_t(((s6[1] >> s_shift) & 0xF) | ((s10[1] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 3)) != 0) ? 0 : 4)), sum)))))))); -+ } -+ temp[j][n] = fma(d, sum, temp[j][n]); -+ } -+ } -+ } -+ -+ reduce_result(temp, d_offset, first_row, num_rows, tid); -+} -+ -+void main() { -+ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); -+ -+ // do NUM_ROWS at a time, unless there aren't enough remaining rows -+ if (first_row + NUM_ROWS <= p.stride_d) { -+ compute_outputs(first_row, NUM_ROWS); -+ } else { -+ if (first_row >= p.stride_d) { -+ return; -+ } -+ compute_outputs(first_row, p.stride_d - first_row); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp -new file mode 100644 -index 00000000..cd1dd8e8 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q4_k.comp -@@ -0,0 +1,133 @@ -+#version 450 -+ -+#extension GL_EXT_shader_explicit_arithmetic_types : require -+ -+#include "mul_mat_vec_base.comp" -+ -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { -+ uint a_offset, b_offset, d_offset; -+ get_offsets(a_offset, b_offset, d_offset); -+ -+ const uint num_blocks_per_row = p.ncols / QUANT_K; -+ -+ // 16 threads are used to process each block -+ const uint it_size = gl_WorkGroupSize.x/16; -+ const uint tid = gl_LocalInvocationID.x; -+ const uint itid = tid%16; // 0...16 -+ const uint ix = tid/16; -+ -+ const uint step = 4; -+ -+ const uint il = itid/step; // 0...3 -+ const uint ir = itid - step*il; // 0...7 or 0...3 -+ const uint n = 4; -+ -+ const uint v_im = il / 2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224 -+ const uint v_in = il % 2; -+ -+ const uint l0 = n * (2 * ir + v_in); // 0...15 -+ const uint q_offset = 32*v_im + l0; -+ const uint y_offset = 64*v_im + l0; -+ -+ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; -+ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { -+ temp[j][i] = FLOAT_TYPE(0); -+ } -+ } -+ -+ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { -+ const uint y1_idx = i * QUANT_K + y_offset; -+ const uint y2_idx = y1_idx + 128; -+ -+ [[unroll]] for (uint n = 0; n < num_rows; ++n) { -+ const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; -+ f16vec2 d = data_a[ib0 + i].d; -+ const FLOAT_TYPE dall = FLOAT_TYPE(d.x); -+ const FLOAT_TYPE dmin = FLOAT_TYPE(d.y); -+ -+ uint32_t scale0_u32 = data_a_packed16[ib0 + i].scales[v_im ]; -+ uint32_t scale4_u32 = data_a_packed16[ib0 + i].scales[v_im + 2]; -+ uint32_t scale8_u32 = data_a_packed16[ib0 + i].scales[v_im + 4]; -+ uvec4 scale0 = uvec4(unpack8(scale0_u32)); -+ uvec4 scale4 = uvec4(unpack8(scale4_u32)); -+ uvec4 scale8 = uvec4(unpack8(scale8_u32)); -+ -+ const uint32_t sc0 = ( scale0.x & 0x3f); -+ const uint32_t sc1 = ( scale0.y & 0x3f); -+ const uint32_t sc2 = ( scale4.x & 0x3f); -+ const uint32_t sc3 = ( scale4.y & 0x3f); -+ const uint32_t sc4 = (( scale8.x & 0x0f) | ((scale0.x & 0xc0) >> 2)); -+ const uint32_t sc5 = (( scale8.y & 0x0f) | ((scale0.y & 0xc0) >> 2)); -+ const uint32_t sc6 = (((scale8.x >> 4) & 0x0f) | ((scale4.x & 0xc0) >> 2)); -+ const uint32_t sc7 = (((scale8.y >> 4) & 0x0f) | ((scale4.y & 0xc0) >> 2)); -+ -+ uint32_t qs0_u32 = data_a_packed32[ib0 + i].qs[q_offset / 4]; -+ uint32_t qs64_u32 = data_a_packed32[ib0 + i].qs[q_offset / 4 + 16]; -+ -+ uint32_t qs0_u32_lo4 = qs0_u32 & 0x0F0F0F0F; -+ uint32_t qs0_u32_hi4 = (qs0_u32 >> 4) & 0x0F0F0F0F; -+ uint32_t qs64_u32_lo4 = qs64_u32 & 0x0F0F0F0F; -+ uint32_t qs64_u32_hi4 = (qs64_u32 >> 4) & 0x0F0F0F0F; -+ -+ uvec4 qs0_lo4 = uvec4(unpack8(qs0_u32_lo4)); -+ uvec4 qs64_lo4 = uvec4(unpack8(qs64_u32_lo4)); -+ uvec4 qs0_hi4 = uvec4(unpack8(qs0_u32_hi4)); -+ uvec4 qs64_hi4 = uvec4(unpack8(qs64_u32_hi4)); -+ -+ const uint32_t q4_0 = qs0_lo4.x; -+ const uint32_t q4_1 = qs0_lo4.y; -+ const uint32_t q4_2 = qs0_lo4.z; -+ const uint32_t q4_3 = qs0_lo4.w; -+ const uint32_t q4_4 = qs0_hi4.x; -+ const uint32_t q4_5 = qs0_hi4.y; -+ const uint32_t q4_6 = qs0_hi4.z; -+ const uint32_t q4_7 = qs0_hi4.w; -+ const uint32_t q4_8 = qs64_lo4.x; -+ const uint32_t q4_9 = qs64_lo4.y; -+ const uint32_t q4_10 = qs64_lo4.z; -+ const uint32_t q4_11 = qs64_lo4.w; -+ const uint32_t q4_12 = qs64_hi4.x; -+ const uint32_t q4_13 = qs64_hi4.y; -+ const uint32_t q4_14 = qs64_hi4.z; -+ const uint32_t q4_15 = qs64_hi4.w; -+ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ B_TYPE_VEC4 by10 = data_b_v4[(j*p.batch_stride_b + b_offset + y1_idx) / 4]; -+ B_TYPE_VEC4 by132 = data_b_v4[(j*p.batch_stride_b + b_offset + y1_idx) / 4 + 8]; -+ B_TYPE_VEC4 by20 = data_b_v4[(j*p.batch_stride_b + b_offset + y2_idx) / 4]; -+ B_TYPE_VEC4 by232 = data_b_v4[(j*p.batch_stride_b + b_offset + y2_idx) / 4 + 8]; -+ -+ const FLOAT_TYPE sx = fma(FLOAT_TYPE(by10.x), q4_0, fma(FLOAT_TYPE(by10.y), q4_1, fma(FLOAT_TYPE(by10.z), q4_2, FLOAT_TYPE(by10.w) * q4_3))); -+ const FLOAT_TYPE sy = fma(FLOAT_TYPE(by132.x), q4_4, fma(FLOAT_TYPE(by132.y), q4_5, fma(FLOAT_TYPE(by132.z), q4_6, FLOAT_TYPE(by132.w) * q4_7))); -+ const FLOAT_TYPE sz = fma(FLOAT_TYPE(by20.x), q4_8, fma(FLOAT_TYPE(by20.y), q4_9, fma(FLOAT_TYPE(by20.z), q4_10, FLOAT_TYPE(by20.w) * q4_11))); -+ const FLOAT_TYPE sw = fma(FLOAT_TYPE(by232.x), q4_12, fma(FLOAT_TYPE(by232.y), q4_13, fma(FLOAT_TYPE(by232.z), q4_14, FLOAT_TYPE(by232.w) * q4_15))); -+ const FLOAT_TYPE smin = -+ fma(FLOAT_TYPE(by10.x), sc2, fma(FLOAT_TYPE(by132.x), sc3, fma(FLOAT_TYPE(by20.x), sc6, fma(FLOAT_TYPE(by232.x), sc7, -+ fma(FLOAT_TYPE(by10.y), sc2, fma(FLOAT_TYPE(by132.y), sc3, fma(FLOAT_TYPE(by20.y), sc6, fma(FLOAT_TYPE(by232.y), sc7, -+ fma(FLOAT_TYPE(by10.z), sc2, fma(FLOAT_TYPE(by132.z), sc3, fma(FLOAT_TYPE(by20.z), sc6, fma(FLOAT_TYPE(by232.z), sc7, -+ fma(FLOAT_TYPE(by10.w), sc2, fma(FLOAT_TYPE(by132.w), sc3, fma(FLOAT_TYPE(by20.w), sc6, FLOAT_TYPE(by232.w) * sc7))))))))))))))); -+ temp[j][n] = fma(dall, fma(sx, sc0, fma(sy, sc1, fma(sz, sc4, sw * sc5))), fma(-dmin, smin, temp[j][n])); -+ } -+ } -+ } -+ -+ reduce_result(temp, d_offset, first_row, num_rows, tid); -+} -+ -+void main() { -+ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); -+ -+ // do NUM_ROWS at a time, unless there aren't enough remaining rows -+ if (first_row + NUM_ROWS <= p.stride_d) { -+ compute_outputs(first_row, NUM_ROWS); -+ } else { -+ if (first_row >= p.stride_d) { -+ return; -+ } -+ compute_outputs(first_row, p.stride_d - first_row); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp -new file mode 100644 -index 00000000..0a68891c ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q5_k.comp -@@ -0,0 +1,162 @@ -+#version 450 -+ -+#extension GL_EXT_shader_explicit_arithmetic_types : require -+ -+#include "mul_mat_vec_base.comp" -+ -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { -+ uint a_offset, b_offset, d_offset; -+ get_offsets(a_offset, b_offset, d_offset); -+ -+ const uint num_blocks_per_row = p.ncols / QUANT_K; -+ -+ // 16 threads are used to process each block -+ const uint it_size = gl_WorkGroupSize.x/16; -+ const uint tid = gl_LocalInvocationID.x; -+ const uint itid = tid%16; // 0...16 -+ const uint ix = tid/16; -+ -+ const uint il = itid/4; // 0...3 -+ const uint ir = itid - 4*il; // 0...7 or 0...3 -+ -+ const uint v_im = il / 2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224 -+ const uint v_in = il % 2; -+ -+ const uint l0 = 4*ir + 2*v_in; // 0...15 -+ const uint q_offset = 32*v_im + l0; -+ const uint y_offset = 64*v_im + l0; -+ -+ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; -+ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { -+ temp[j][i] = FLOAT_TYPE(0); -+ } -+ } -+ -+ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { -+ const uint y1_idx = i * QUANT_K + y_offset; -+ const uint y2_idx = y1_idx + 128; -+ -+ [[unroll]] for (uint n = 0; n < num_rows; ++n) { -+ const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; -+ f16vec2 d = data_a[ib0 + i].d; -+ const FLOAT_TYPE dall = FLOAT_TYPE(d.x); -+ const FLOAT_TYPE dmin = FLOAT_TYPE(d.y); -+ -+ uint32_t scale0_u32 = data_a_packed16[ib0 + i].scales[v_im ]; -+ uint32_t scale4_u32 = data_a_packed16[ib0 + i].scales[v_im + 2]; -+ uint32_t scale8_u32 = data_a_packed16[ib0 + i].scales[v_im + 4]; -+ uvec4 scale0 = uvec4(unpack8(scale0_u32)); -+ uvec4 scale4 = uvec4(unpack8(scale4_u32)); -+ uvec4 scale8 = uvec4(unpack8(scale8_u32)); -+ -+ const uint32_t sc0 = ( scale0.x & 0x3f); -+ const uint32_t sc1 = ( scale0.y & 0x3f); -+ const uint32_t sc2 = ( scale4.x & 0x3f); -+ const uint32_t sc3 = ( scale4.y & 0x3f); -+ const uint32_t sc4 = (( scale8.x & 0x0f) | ((scale0.x & 0xc0) >> 2)); -+ const uint32_t sc5 = (( scale8.y & 0x0f) | ((scale0.y & 0xc0) >> 2)); -+ const uint32_t sc6 = (((scale8.x >> 4) & 0x0f) | ((scale4.x & 0xc0) >> 2)); -+ const uint32_t sc7 = (((scale8.y >> 4) & 0x0f) | ((scale4.y & 0xc0) >> 2)); -+ -+ uint32_t qs0_16_u32 = uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 8]) << 16); -+ uint32_t qs64_80_u32 = uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 32]) | (uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 40]) << 16); -+ -+ uint32_t qs0_16_u32_lo4 = qs0_16_u32 & 0x0F0F0F0F; -+ uint32_t qs0_16_u32_hi4 = (qs0_16_u32 >> 4) & 0x0F0F0F0F; -+ uint32_t qs64_80_u32_lo4 = qs64_80_u32 & 0x0F0F0F0F; -+ uint32_t qs64_80_u32_hi4 = (qs64_80_u32 >> 4) & 0x0F0F0F0F; -+ -+ uint32_t qh = pack32(u16vec2(data_a_packed16[ib0 + i].qh[l0 / 2], data_a_packed16[ib0 + i].qh[l0 / 2 + 8])); -+ -+ uint32_t qs0_16_lo4_offset16 = ((qh >> (2*v_im)) & 0x01010101) << 4; -+ uint32_t qs0_16_hi4_offset16 = ((qh >> (2*v_im)) & 0x02020202) << 3; -+ uint32_t qs64_80_lo4_offset16 = ((qh >> (2*v_im)) & 0x10101010) << 0; -+ uint32_t qs64_80_hi4_offset16 = ((qh >> (2*v_im)) & 0x20202020) >> 1; -+ -+ qs0_16_u32_lo4 += qs0_16_lo4_offset16; -+ qs0_16_u32_hi4 += qs0_16_hi4_offset16; -+ qs64_80_u32_lo4 += qs64_80_lo4_offset16; -+ qs64_80_u32_hi4 += qs64_80_hi4_offset16; -+ -+ uvec4 qs0_16_lo4 = uvec4(unpack8(qs0_16_u32_lo4)); -+ uvec4 qs64_80_lo4 = uvec4(unpack8(qs64_80_u32_lo4)); -+ uvec4 qs0_16_hi4 = uvec4(unpack8(qs0_16_u32_hi4)); -+ uvec4 qs64_80_hi4 = uvec4(unpack8(qs64_80_u32_hi4)); -+ -+ const uint32_t q4_0 = qs0_16_lo4.x; -+ const uint32_t q4_1 = qs0_16_lo4.y; -+ const uint32_t q4_2 = qs0_16_lo4.z; -+ const uint32_t q4_3 = qs0_16_lo4.w; -+ const uint32_t q4_4 = qs0_16_hi4.x; -+ const uint32_t q4_5 = qs0_16_hi4.y; -+ const uint32_t q4_6 = qs0_16_hi4.z; -+ const uint32_t q4_7 = qs0_16_hi4.w; -+ const uint32_t q4_8 = qs64_80_lo4.x; -+ const uint32_t q4_9 = qs64_80_lo4.y; -+ const uint32_t q4_10 = qs64_80_lo4.z; -+ const uint32_t q4_11 = qs64_80_lo4.w; -+ const uint32_t q4_12 = qs64_80_hi4.x; -+ const uint32_t q4_13 = qs64_80_hi4.y; -+ const uint32_t q4_14 = qs64_80_hi4.z; -+ const uint32_t q4_15 = qs64_80_hi4.w; -+ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ B_TYPE_VEC2 by10 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2]; -+ B_TYPE_VEC2 by116 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 8]; -+ B_TYPE_VEC2 by132 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 16]; -+ B_TYPE_VEC2 by148 = data_b_v2[(j*p.batch_stride_b + b_offset + y1_idx) / 2 + 24]; -+ B_TYPE_VEC2 by20 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2]; -+ B_TYPE_VEC2 by216 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 8]; -+ B_TYPE_VEC2 by232 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 16]; -+ B_TYPE_VEC2 by248 = data_b_v2[(j*p.batch_stride_b + b_offset + y2_idx) / 2 + 24]; -+ -+ const FLOAT_TYPE sx = -+ fma(FLOAT_TYPE(by10.x), q4_0, -+ fma(FLOAT_TYPE(by10.y), q4_1, -+ fma(FLOAT_TYPE(by116.x), q4_2, -+ FLOAT_TYPE(by116.y) * q4_3))); -+ const FLOAT_TYPE sy = -+ fma(FLOAT_TYPE(by132.x), q4_4, -+ fma(FLOAT_TYPE(by132.y), q4_5, -+ fma(FLOAT_TYPE(by148.x), q4_6, -+ FLOAT_TYPE(by148.y) * q4_7))); -+ const FLOAT_TYPE sz = -+ fma(FLOAT_TYPE(by20.x), q4_8, -+ fma(FLOAT_TYPE(by20.y), q4_9, -+ fma(FLOAT_TYPE(by216.x), q4_10, -+ FLOAT_TYPE(by216.y) * q4_11))); -+ const FLOAT_TYPE sw = -+ fma(FLOAT_TYPE(by232.x), q4_12, -+ fma(FLOAT_TYPE(by232.y), q4_13, -+ fma(FLOAT_TYPE(by248.x), q4_14, -+ FLOAT_TYPE(by248.y) * q4_15))); -+ const FLOAT_TYPE smin = -+ fma(FLOAT_TYPE(by10.x) + FLOAT_TYPE(by10.y) + FLOAT_TYPE(by116.x) + FLOAT_TYPE(by116.y), sc2, -+ fma(FLOAT_TYPE(by132.x) + FLOAT_TYPE(by132.y) + FLOAT_TYPE(by148.x) + FLOAT_TYPE(by148.y), sc3, -+ fma(FLOAT_TYPE(by20.x) + FLOAT_TYPE(by20.y) + FLOAT_TYPE(by216.x) + FLOAT_TYPE(by216.y), sc6, -+ (FLOAT_TYPE(by232.x) + FLOAT_TYPE(by232.y) + FLOAT_TYPE(by248.x) + FLOAT_TYPE(by248.y)) * sc7))); -+ temp[j][n] = fma(dall, fma(sx, sc0, fma(sy, sc1, fma(sz, sc4, sw * sc5))), fma(-dmin, smin, temp[j][n])); -+ } -+ } -+ } -+ -+ reduce_result(temp, d_offset, first_row, num_rows, tid); -+} -+ -+void main() { -+ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); -+ -+ // do NUM_ROWS at a time, unless there aren't enough remaining rows -+ if (first_row + NUM_ROWS <= p.stride_d) { -+ compute_outputs(first_row, NUM_ROWS); -+ } else { -+ if (first_row >= p.stride_d) { -+ return; -+ } -+ compute_outputs(first_row, p.stride_d - first_row); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp -new file mode 100644 -index 00000000..70e13a56 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q6_k.comp -@@ -0,0 +1,112 @@ -+#version 450 -+ -+#extension GL_EXT_shader_explicit_arithmetic_types : require -+ -+#include "mul_mat_vec_base.comp" -+ -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { -+ uint a_offset, b_offset, d_offset; -+ get_offsets(a_offset, b_offset, d_offset); -+ -+ const uint num_blocks_per_row = p.ncols / QUANT_K; -+ -+ // 16 threads are used to process each block -+ const uint it_size = gl_WorkGroupSize.x/16; -+ const uint tid = gl_LocalInvocationID.x; -+ const uint itid = tid%16; // 0...16 -+ const uint ix = tid/16; -+ -+ const uint step = 8; -+ -+ const uint v_im = itid/step; // 0 or 1. 0 computes 0..., 1 computes 128... -+ const uint v_in = itid - step*v_im; // 0...15 or 0...7 -+ -+ const uint l0 = 4 * v_in; // 0, 4, 8, ..., 28 -+ const uint is = v_in / 4; -+ -+ const uint ql_offset = 64*v_im + l0; -+ const uint qh_offset = 32*v_im + l0; -+ const uint s_offset = 8*v_im + is; -+ const uint y_offset = 128*v_im + l0; -+ -+ FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; -+ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { -+ temp[j][i] = FLOAT_TYPE(0); -+ } -+ } -+ -+ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) { -+ const uint y_idx = i * QUANT_K + y_offset; -+ -+ [[unroll]] for (uint n = 0; n < num_rows; ++n) { -+ const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row; -+ const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d); -+ -+ FLOAT_TYPE scales[4]; -+ scales[0] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 0]); -+ scales[1] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 2]); -+ scales[2] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 4]); -+ scales[3] = FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 6]); -+ -+ uint32_t ql0_u32 = uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 1]) << 16); -+ uint32_t ql32_u32 = uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 16]) | (uint32_t(data_a_packed16[ib0 + i].ql[ql_offset / 2 + 17]) << 16); -+ -+ uint32_t ql0_u32_lo4 = ql0_u32 & 0x0F0F0F0F; -+ uint32_t ql0_u32_hi4 = (ql0_u32 >> 4) & 0x0F0F0F0F; -+ uint32_t ql32_u32_lo4 = ql32_u32 & 0x0F0F0F0F; -+ uint32_t ql32_u32_hi4 = (ql32_u32 >> 4) & 0x0F0F0F0F; -+ -+ uint32_t qh_u32 = uint32_t(data_a_packed16[ib0 + i].qh[qh_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].qh[qh_offset / 2 + 1]) << 16); -+ uint32_t qh0_u32 = (qh_u32 & 0x03030303) << 4; -+ uint32_t qh2_u32 = (qh_u32 & 0x0C0C0C0C) << 2; -+ uint32_t qh4_u32 = (qh_u32 & 0x30303030) << 0; -+ uint32_t qh6_u32 = (qh_u32 & 0xC0C0C0C0) >> 2; -+ -+ uint32_t q0_u32 = ql0_u32_lo4 | qh0_u32; -+ uint32_t q1_u32 = ql32_u32_lo4 | qh2_u32; -+ uint32_t q2_u32 = ql0_u32_hi4 | qh4_u32; -+ uint32_t q3_u32 = ql32_u32_hi4 | qh6_u32; -+ -+ uvec4 q0 = uvec4(unpack8(q0_u32)); -+ uvec4 q1 = uvec4(unpack8(q1_u32)); -+ uvec4 q2 = uvec4(unpack8(q2_u32)); -+ uvec4 q3 = uvec4(unpack8(q3_u32)); -+ -+ [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { -+ B_TYPE_VEC4 by0 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4]; -+ B_TYPE_VEC4 by32 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 8]; -+ B_TYPE_VEC4 by64 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 16]; -+ B_TYPE_VEC4 by96 = data_b_v4[(j*p.batch_stride_b + b_offset + y_idx) / 4 + 24]; -+ -+ FLOAT_TYPE sum = FLOAT_TYPE(0.0); -+ [[unroll]] for (int l = 0; l < 4; ++l) { -+ sum = fma(FLOAT_TYPE(by0[l]) * scales[0], FLOAT_TYPE(int8_t(q0[l]) - 32), -+ fma(FLOAT_TYPE(by32[l]) * scales[1], FLOAT_TYPE(int8_t(q1[l]) - 32), -+ fma(FLOAT_TYPE(by64[l]) * scales[2], FLOAT_TYPE(int8_t(q2[l]) - 32), -+ fma(FLOAT_TYPE(by96[l]) * scales[3], FLOAT_TYPE(int8_t(q3[l]) - 32), sum)))); -+ } -+ temp[j][n] += sum * d; -+ } -+ } -+ } -+ -+ reduce_result(temp, d_offset, first_row, num_rows, tid); -+} -+ -+void main() { -+ const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); -+ -+ // do NUM_ROWS at a time, unless there aren't enough remaining rows -+ if (first_row + NUM_ROWS <= p.stride_d) { -+ compute_outputs(first_row, NUM_ROWS); -+ } else { -+ if (first_row >= p.stride_d) { -+ return; -+ } -+ compute_outputs(first_row, p.stride_d - first_row); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp -new file mode 100644 -index 00000000..48122cbe ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp -@@ -0,0 +1,631 @@ -+#version 450 -+ -+#extension GL_EXT_control_flow_attributes : enable -+#extension GL_EXT_shader_16bit_storage : require -+ -+#ifdef FLOAT16 -+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -+#endif -+ -+#ifdef COOPMAT -+#extension GL_KHR_cooperative_matrix : enable -+#extension GL_KHR_memory_scope_semantics : enable -+#extension GL_KHR_shader_subgroup_basic : enable -+#endif -+ -+#ifdef MUL_MAT_ID -+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require -+#endif -+ -+#include "types.comp" -+ -+#ifndef LOAD_VEC_A -+#define LOAD_VEC_A 1 -+#endif -+#ifndef LOAD_VEC_B -+#define LOAD_VEC_B 1 -+#endif -+ -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; -+layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; -+ -+#ifdef MUL_MAT_ID -+layout (binding = 3) readonly buffer IDS {int data_ids[];}; -+#endif -+ -+layout (push_constant) uniform parameter -+{ -+ uint M; -+ uint N; -+ uint K; -+ uint stride_a; -+ uint stride_b; -+ uint stride_d; -+ -+ uint batch_stride_a; -+ uint batch_stride_b; -+ uint batch_stride_d; -+ -+#ifdef MUL_MAT_ID -+ uint nei0; -+ uint nei1; -+ uint nbi1; -+ uint ne11; -+#else -+ uint k_split; -+ uint ne02; -+ uint ne12; -+ uint broadcast2; -+ uint broadcast3; -+#endif -+} p; -+ -+layout (constant_id = 0) const uint BLOCK_SIZE = 64; -+layout (constant_id = 1) const uint BM = 64; -+layout (constant_id = 2) const uint BN = 64; -+layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working with a quant -+layout (constant_id = 4) const uint WM = 32; -+layout (constant_id = 5) const uint WN = 32; -+layout (constant_id = 6) const uint WMITER = 2; -+layout (constant_id = 7) const uint TM = 4; -+layout (constant_id = 8) const uint TN = 2; -+layout (constant_id = 9) const uint TK = 1; // Only needed for coopmat -+layout (constant_id = 10) const uint WARP = 32; -+ -+#ifdef COOPMAT -+#define SHMEM_STRIDE (BK + 8) -+#else -+#define SHMEM_STRIDE (BK + 1) -+#endif -+ -+shared FLOAT_TYPE buf_a[BM * SHMEM_STRIDE]; -+shared FLOAT_TYPE buf_b[BN * SHMEM_STRIDE]; -+ -+#ifdef MUL_MAT_ID -+shared u16vec2 row_ids[3072]; -+#endif // MUL_MAT_ID -+ -+#define NUM_WARPS (BLOCK_SIZE / WARP) -+ -+#ifdef COOPMAT -+shared ACC_TYPE coopmat_stage[TM * TN * NUM_WARPS]; -+#endif -+ -+void main() { -+#if defined(DATA_A_IQ4_NL) -+ init_iq4nl_shmem(); -+#endif -+ -+#ifdef MUL_MAT_ID -+ const uint expert_idx = gl_GlobalInvocationID.z; -+#else -+ const uint batch_idx = gl_GlobalInvocationID.z; -+ -+ const uint i13 = batch_idx / p.ne12; -+ const uint i12 = batch_idx % p.ne12; -+ -+ const uint i03 = i13 / p.broadcast3; -+ const uint i02 = i12 / p.broadcast2; -+ -+ const uint batch_idx_a = i03 * p.ne02 + i02; -+#endif -+ -+ const uint blocks_m = (p.M + BM - 1) / BM; -+ const uint ir = gl_WorkGroupID.x % blocks_m; -+ const uint ik = gl_WorkGroupID.x / blocks_m; -+ const uint ic = gl_WorkGroupID.y; -+ -+ const uint WNITER = (WM * WN) / (WARP * TM * TN * WMITER); -+ const uint WSUBM = WM / WMITER; -+ const uint WSUBN = WN / WNITER; -+ -+#ifdef COOPMAT -+ const uint warp_i = gl_SubgroupID; -+ -+ const uint tiw = gl_SubgroupInvocationID; -+ -+ const uint cms_per_row = WM / TM; -+ const uint cms_per_col = WN / TN; -+ -+ const uint storestride = WARP / TM; -+ const uint store_r = tiw % TM; -+ const uint store_c = tiw / TM; -+#else -+ const uint warp_i = gl_LocalInvocationID.x / WARP; -+ -+ const uint tiw = gl_LocalInvocationID.x % WARP; -+ -+ const uint tiwr = tiw % (WSUBM / TM); -+ const uint tiwc = tiw / (WSUBM / TM); -+#endif -+ -+ const uint warp_r = warp_i % (BM / WM); -+ const uint warp_c = warp_i / (BM / WM); -+ -+ const uint loadr_a = gl_LocalInvocationID.x % (BK / LOAD_VEC_A); -+ const uint loadc_a = gl_LocalInvocationID.x / (BK / LOAD_VEC_A); -+ const uint loadr_b = gl_LocalInvocationID.x % (BK / LOAD_VEC_B); -+ const uint loadc_b = gl_LocalInvocationID.x / (BK / LOAD_VEC_B); -+ -+ const uint loadstride_a = gl_WorkGroupSize.x * LOAD_VEC_A / BK; -+ const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B / BK; -+ -+#ifdef MUL_MAT_ID -+ uint _ne1 = 0; -+ for (uint ii1 = 0; ii1 < p.nei1; ii1++) { -+ for (uint ii0 = 0; ii0 < p.nei0; ii0++) { -+ if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { -+ row_ids[_ne1] = u16vec2(ii0, ii1); -+ _ne1++; -+ } -+ } -+ } -+ -+ barrier(); -+ -+ // Workgroup has no work -+ if (ic * BN >= _ne1) return; -+#endif -+ -+#ifdef MUL_MAT_ID -+ const uint start_k = 0; -+ const uint end_k = p.K; -+#else -+ const uint start_k = ik * p.k_split; -+ const uint end_k = min(p.K, (ik + 1) * p.k_split); -+#endif -+ -+ uint pos_a = ( -+#ifdef MUL_MAT_ID -+ expert_idx * p.batch_stride_a + -+#else -+ batch_idx_a * p.batch_stride_a + -+#endif -+ ir * BM * p.stride_a + start_k) / LOAD_VEC_A; -+#ifdef MUL_MAT_ID -+ uint pos_b = 0; -+#else -+ uint pos_b = (batch_idx * p.batch_stride_b + ic * BN * p.stride_b + start_k) / LOAD_VEC_B; -+#endif -+ -+#ifdef COOPMAT -+ coopmat cache_a; -+ coopmat cache_b; -+ coopmat sums[cms_per_row * cms_per_col]; -+ -+ [[unroll]] for (uint i = 0; i < cms_per_row * cms_per_col; i++) { -+ sums[i] = coopmat(0.0f); -+ } -+#else -+ ACC_TYPE sums[WMITER * TM * WNITER * TN]; -+ FLOAT_TYPE cache_a[WMITER * TM]; -+ FLOAT_TYPE cache_b[WNITER * TN]; -+ -+ [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN; i++) { -+ sums[i] = ACC_TYPE(0.0f); -+ } -+#endif -+ -+ for (uint block = start_k; block < end_k; block += BK) { -+ [[unroll]] for (uint l = 0; l < BM; l += loadstride_a) { -+ -+#if defined(DATA_A_F32) || defined(DATA_A_F16) -+#if LOAD_VEC_A == 8 -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; -+ buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx][0].x); -+ buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx][0].y); -+ buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx][0].z); -+ buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx][0].w); -+ buf_a[buf_idx + 4] = FLOAT_TYPE(data_a[idx][1].x); -+ buf_a[buf_idx + 5] = FLOAT_TYPE(data_a[idx][1].y); -+ buf_a[buf_idx + 6] = FLOAT_TYPE(data_a[idx][1].z); -+ buf_a[buf_idx + 7] = FLOAT_TYPE(data_a[idx][1].w); -+#elif LOAD_VEC_A == 4 -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; -+ buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx].x); -+ buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx].y); -+ buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx].z); -+ buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx].w); -+#else -+ if (ir * BM + loadc_a + l < p.M && block + loadr_a < end_k) { -+ buf_a[(loadc_a + l) * SHMEM_STRIDE + loadr_a] = FLOAT_TYPE(data_a[pos_a + (loadc_a + l) * p.stride_a + loadr_a]); -+ } else { -+ buf_a[(loadc_a + l) * SHMEM_STRIDE + loadr_a] = FLOAT_TYPE(0.0f); -+ } -+#endif -+#elif defined(DATA_A_Q4_0) -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; -+ -+ const uint ib = idx / 16; -+ const uint iqs = idx & 0xF; -+ -+ const float d = float(data_a[ib].d); -+ const uint vui = uint(data_a[ib].qs[iqs]); -+ const vec2 v = (vec2(vui & 0xF, vui >> 4) - 8.0f) * d; -+ -+ buf_a[buf_idx ] = FLOAT_TYPE(v.x); -+ buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); -+#elif defined(DATA_A_Q4_1) -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; -+ -+ const uint ib = idx / 16; -+ const uint iqs = idx & 0xF; -+ -+ const float d = float(data_a[ib].d); -+ const float m = float(data_a[ib].m); -+ const uint vui = uint(data_a[ib].qs[iqs]); -+ const vec2 v = vec2(vui & 0xF, vui >> 4) * d + m; -+ -+ buf_a[buf_idx ] = FLOAT_TYPE(v.x); -+ buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); -+#elif defined(DATA_A_Q5_0) -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; -+ -+ const uint ib = idx / 16; -+ const uint iqs = idx & 0xF; -+ -+ const float d = float(data_a[ib].d); -+ const uint uint_qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0]; -+ const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); -+ const uint vui = uint(data_a[ib].qs[iqs]); -+ const vec2 v = (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f) * d; -+ -+ buf_a[buf_idx ] = FLOAT_TYPE(v.x); -+ buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); -+#elif defined(DATA_A_Q5_1) -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; -+ -+ const uint ib = idx / 16; -+ const uint iqs = idx & 0xF; -+ -+ const float d = float(data_a[ib].d); -+ const float m = float(data_a[ib].m); -+ const uint uint_qh = data_a[ib].qh; -+ const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); -+ const uint vui = uint(data_a[ib].qs[iqs]); -+ const vec2 v = vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) * d + m; -+ -+ buf_a[buf_idx ] = FLOAT_TYPE(v.x); -+ buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); -+#elif defined(DATA_A_Q8_0) -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; -+ -+ const uint ib = idx / 16; -+ const uint iqs = (idx & 0xF) * 2; -+ -+ const float d = float(data_a[ib].d); -+ const vec2 v = vec2(int(data_a[ib].qs[iqs]), int(data_a[ib].qs[iqs + 1])) * d; -+ -+ buf_a[buf_idx ] = FLOAT_TYPE(v.x); -+ buf_a[buf_idx + 1] = FLOAT_TYPE(v.y); -+#elif defined(DATA_A_Q2_K) -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; -+ -+ const uint ib = idx / 128; // 2 values per idx -+ const uint iqs = idx % 128; // 0..127 -+ -+ const uint qsi = (iqs / 64) * 32 + (iqs % 16) * 2; // 0,2,4..30 -+ const uint scalesi = iqs / 8; // 0..15 -+ const uint qsshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 -+ -+ const uvec2 qs = uvec2(data_a[ib].qs[qsi], data_a[ib].qs[qsi + 1]); -+ const uint scales = data_a[ib].scales[scalesi]; -+ const vec2 d = vec2(data_a[ib].d); -+ -+ const vec2 v = d.x * float(scales & 0xF) * vec2((qs >> qsshift) & 3) - d.y * float(scales >> 4); -+ -+ buf_a[buf_idx ] = FLOAT_TYPE(v.x); -+ buf_a[buf_idx + 1] = FLOAT_TYPE(v.y); -+#elif defined(DATA_A_Q3_K) -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; -+ -+ const uint ib = idx / 128; // 2 values per idx -+ const uint iqs = idx % 128; // 0..127 -+ -+ const uint n = iqs / 64; // 0,1 -+ const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 -+ const uint hmi = (iqs % 16) * 2; // 0,2,4..30 -+ const uint j = (iqs % 64) / 4; // 0..3 -+ const uint is = iqs / 8; // 0..15 -+ const uint halfsplit = ((iqs % 64) / 16); // 0,1,2,3 -+ const uint qsshift = halfsplit * 2; // 0,2,4,6 -+ const uint m = 1 << (4 * n + halfsplit); // 1,2,4,8,16,32,64,128 -+ -+ const int8_t us = int8_t(is < 4 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+8] >> 0) & 3) << 4) : -+ is < 8 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+4] >> 2) & 3) << 4) : -+ is < 12 ? (data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is+0] >> 4) & 3) << 4) : -+ (data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is-4] >> 6) & 3) << 4)); -+ const float dl = float(data_a[ib].d) * float(us - 32); -+ -+ buf_a[buf_idx ] = FLOAT_TYPE(dl * float(int8_t((data_a[ib].qs[qsi ] >> qsshift) & 3) - (((data_a[ib].hmask[hmi ] & m) != 0) ? 0 : 4))); -+ buf_a[buf_idx + 1] = FLOAT_TYPE(dl * float(int8_t((data_a[ib].qs[qsi + 1] >> qsshift) & 3) - (((data_a[ib].hmask[hmi + 1] & m) != 0) ? 0 : 4))); -+#elif defined(DATA_A_Q4_K) -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; -+ -+ const uint ib = idx / 128; // 2 values per idx -+ const uint iqs = idx % 128; // 0..127 -+ -+ const uint n = iqs / 32; // 0,1,2,3 -+ const uint b = (iqs % 32) / 16; // 0,1 -+ const uint is = 2 * n + b; // 0..7 -+ const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 -+ -+ const vec2 loadd = vec2(data_a[ib].d); -+ -+ const uint scidx0 = (is < 4) ? is : (is + 4); -+ const uint scidx1 = (is < 4) ? is : (is - 4); -+ const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ const uint scidxshift1 = (is < 4) ? 0 : 2; -+ const uint mbidx0 = is + 4; -+ const uint mbidx1 = (is < 4) ? is + 4 : is; -+ const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; -+ const uint mbidxshift0 = (is < 4) ? 0 : 4; -+ const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ const uint mbidxshift1 = (is < 4) ? 0 : 2; -+ -+ const uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); -+ const uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); -+ -+ const float d = loadd.x * sc; -+ const float m = -loadd.y * mbyte; -+ -+ buf_a[buf_idx ] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi ] >> (b * 4)) & 0xF), m)); -+ buf_a[buf_idx + 1] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi + 1] >> (b * 4)) & 0xF), m)); -+#elif defined(DATA_A_Q5_K) -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; -+ -+ const uint ib = idx / 128; // 2 values per idx -+ const uint iqs = idx % 128; // 0..127 -+ -+ const uint n = iqs / 32; // 0,1,2,3 -+ const uint b = (iqs % 32) / 16; // 0,1 -+ const uint is = 2 * n + b; // 0..7 -+ const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 -+ const uint qhi = (iqs % 16) * 2; // 0,2,4..30 -+ -+ const uint8_t hm = uint8_t(1 << (iqs / 16)); -+ -+ const vec2 loadd = vec2(data_a[ib].d); -+ -+ const uint scidx0 = (is < 4) ? is : (is + 4); -+ const uint scidx1 = (is < 4) ? is : (is - 4); -+ const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ const uint scidxshift1 = (is < 4) ? 0 : 2; -+ const uint mbidx0 = is + 4; -+ const uint mbidx1 = (is < 4) ? is + 4 : is; -+ const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0; -+ const uint mbidxshift0 = (is < 4) ? 0 : 4; -+ const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0; -+ const uint mbidxshift1 = (is < 4) ? 0 : 2; -+ -+ const uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1)); -+ const uint8_t mbyte = uint8_t(((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1)); -+ -+ const float d = loadd.x * sc; -+ const float m = -loadd.y * mbyte; -+ -+ buf_a[buf_idx ] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi ] >> (b * 4)) & 0xF) + float((data_a[ib].qh[qhi ] & hm) != 0 ? 16 : 0), m)); -+ buf_a[buf_idx + 1] = FLOAT_TYPE(fma(d, float((data_a[ib].qs[qsi + 1] >> (b * 4)) & 0xF) + float((data_a[ib].qh[qhi + 1] & hm) != 0 ? 16 : 0), m)); -+#elif defined(DATA_A_Q6_K) -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A; -+ -+ const uint ib = idx / 128; // 2 values per idx -+ const uint iqs = idx % 128; // 0..127 -+ -+ const uint n = iqs / 64; // 0,1 -+ const uint b = (iqs % 64) / 32; // 0,1 -+ const uint is_b = (iqs % 16) / 8; // 0,1 -+ const uint qhshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 -+ const uint is = 8 * n + qhshift + is_b; // 0..15 -+ const uint qsi = n * 64 + (iqs % 32) * 2; // 0,2,4..126 -+ const uint qhi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 -+ -+ const float dscale = float(data_a[ib].d) * float(data_a[ib].scales[is]); -+ -+ buf_a[buf_idx ] = FLOAT_TYPE(dscale * float(int8_t(((data_a[ib].ql[qsi ] >> (b * 4)) & 0xF) | (((data_a[ib].qh[qhi ] >> qhshift) & 3) << 4)) - 32)); -+ buf_a[buf_idx + 1] = FLOAT_TYPE(dscale * float(int8_t(((data_a[ib].ql[qsi + 1] >> (b * 4)) & 0xF) | (((data_a[ib].qh[qhi + 1] >> qhshift) & 3) << 4)) - 32)); -+#elif defined(DATA_A_IQ4_NL) -+ const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; -+ const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a; -+ -+ const uint ib = idx / 16; -+ const uint iqs = idx & 0xF; -+ -+ const float d = float(data_a[ib].d); -+ const uint vui = uint(data_a[ib].qs[iqs]); -+ const vec2 v = vec2(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[vui >> 4]) * d; -+ -+ buf_a[buf_idx ] = FLOAT_TYPE(v.x); -+ buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); -+#endif -+ } -+ [[unroll]] for (uint l = 0; l < BN; l += loadstride_b) { -+#if LOAD_VEC_B == 8 -+#ifdef MUL_MAT_ID -+ const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l]; -+ const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + loadr_b; -+#else -+ const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b; -+#endif -+ const uint buf_idx = (loadc_b + l) * SHMEM_STRIDE + loadr_b * LOAD_VEC_B; -+ buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx][0].x); -+ buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx][0].y); -+ buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx][0].z); -+ buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx][0].w); -+ buf_b[buf_idx + 4] = FLOAT_TYPE(data_b[idx][1].x); -+ buf_b[buf_idx + 5] = FLOAT_TYPE(data_b[idx][1].y); -+ buf_b[buf_idx + 6] = FLOAT_TYPE(data_b[idx][1].z); -+ buf_b[buf_idx + 7] = FLOAT_TYPE(data_b[idx][1].w); -+#elif LOAD_VEC_B == 4 -+#ifdef MUL_MAT_ID -+ const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l]; -+ const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + loadr_b; -+#else -+ const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b; -+#endif -+ const uint buf_idx = (loadc_b + l) * SHMEM_STRIDE + loadr_b * LOAD_VEC_B; -+ buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx].x); -+ buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx].y); -+ buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx].z); -+ buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx].w); -+#elif !MUL_MAT_ID -+ if (ic * BN + loadc_b + l < p.N && block + loadr_b < end_k) { -+ buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(data_b[pos_b + (loadc_b + l) * p.stride_b + loadr_b]); -+ } else { -+ buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(0.0f); -+ } -+#else -+ const uint row_i = ic * BN + loadc_b + l; -+ if (row_i < _ne1) { -+ const u16vec2 row_idx = row_ids[row_i]; -+ buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(data_b[pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + loadr_b]); -+ } else { -+ buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = FLOAT_TYPE(0.0f); -+ } -+#endif -+ } -+ -+ barrier(); -+ -+ pos_a += BK / LOAD_VEC_A; -+ pos_b += BK / LOAD_VEC_B; -+ -+#ifdef COOPMAT -+ [[unroll]] for (uint i = 0; i < BK; i += TK) { -+ [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { -+ // Load from shared into cache -+ coopMatLoad(cache_a, buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor); -+ -+ [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { -+ coopMatLoad(cache_b, buf_b, (warp_c * WN + cm_col * TN) * SHMEM_STRIDE + i, SHMEM_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); -+ -+ sums[cm_col * cms_per_row + cm_row] = coopMatMulAdd(cache_a, cache_b, sums[cm_col * cms_per_row + cm_row]); -+ } -+ } -+ } -+#else -+ [[unroll]] for (uint i = 0; i < BK; i++) { -+ // Load from shared into cache -+ [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { -+ [[unroll]] for (uint j = 0; j < TM; j++) { -+ cache_a[wsir * TM + j] = buf_a[(warp_r * WM + wsir * WSUBM + tiwr * TM + j) * SHMEM_STRIDE + i]; -+ } -+ } -+ [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { -+ [[unroll]] for (uint j = 0; j < TN; j++) { -+ cache_b[wsic * TN + j] = buf_b[(warp_c * WN + wsic * WSUBN + tiwc * TN + j) * SHMEM_STRIDE + i]; -+ } -+ } -+ -+ [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { -+ [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { -+ [[unroll]] for (uint cc = 0; cc < TN; cc++) { -+ [[unroll]] for (uint cr = 0; cr < TM; cr++) { -+ const uint sums_idx = (wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr; -+ sums[sums_idx] = fma(ACC_TYPE(cache_a[wsir * TM + cr]), ACC_TYPE(cache_b[wsic * TN + cc]), sums[sums_idx]); -+ } -+ } -+ } -+ } -+ } -+#endif -+ -+ barrier(); -+ } -+ -+ const uint dr = ir * BM + warp_r * WM; -+ const uint dc = ic * BN + warp_c * WN; -+ -+#ifndef MUL_MAT_ID -+ const uint offsets = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * gl_NumWorkGroups.z; -+#endif -+ -+#ifdef COOPMAT -+#ifdef MUL_MAT_ID -+ [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { -+ [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { -+ coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); -+ -+ [[unroll]] for (uint col = 0; col < BN; col += storestride) { -+ const uint row_i = dc + cm_col * TN + col + store_c; -+ if (row_i >= _ne1) break; -+ -+ const u16vec2 row_idx = row_ids[row_i]; -+ -+ data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); -+ } -+ } -+ } -+#else -+ const bool is_aligned = p.stride_d % 4 == 0; // Assumption: D_TYPE == float -+ -+ [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { -+ [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { -+ const bool is_in_bounds = dr + (cm_row + 1) * TM <= p.M && dc + (cm_col + 1) * TN <= p.N; -+ -+ if (is_aligned && is_in_bounds) { -+ // Full coopMat is within bounds and stride_d is aligned with 16B -+ coopmat cm_dtype = coopmat(sums[cm_col * cms_per_row + cm_row]); -+ coopMatStore(cm_dtype, data_d, offsets + (dc + cm_col * TN) * p.stride_d + dr + cm_row * TM, p.stride_d, gl_CooperativeMatrixLayoutColumnMajor); -+ } else if (is_in_bounds) { -+ // Full coopMat is within bounds, but stride_d is not aligned -+ coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); -+ -+ [[unroll]] for (uint col = 0; col < TN; col += storestride) { -+ data_d[offsets + (dc + cm_col * TN + col + store_c) * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); -+ } -+ } else if (dr + cm_row * TM < p.M && dc + cm_col * TN < p.N) { -+ // Partial coopMat is within bounds -+ coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor); -+ -+ [[unroll]] for (uint col = 0; col < TN; col += storestride) { -+ if (dr + cm_row * TM + store_r < p.M && dc + cm_col * TN + col + store_c < p.N) { -+ data_d[offsets + (dc + cm_col * TN + col + store_c) * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); -+ } -+ } -+ } -+ } -+ } -+#endif // MUL_MAT_ID -+#else -+ [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { -+ [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { -+ -+ const uint dr_warp = dr + wsir * WSUBM + tiwr * TM; -+ const uint dc_warp = dc + wsic * WSUBN + tiwc * TN; -+ [[unroll]] for (uint cc = 0; cc < TN; cc++) { -+#ifdef MUL_MAT_ID -+ const uint row_i = dc_warp + cc; -+ if (row_i >= _ne1) break; -+ -+ const u16vec2 row_idx = row_ids[row_i]; -+#endif // MUL_MAT_ID -+ [[unroll]] for (uint cr = 0; cr < TM; cr++) { -+#ifdef MUL_MAT_ID -+ data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]); -+#else -+ if (dr_warp + cr < p.M && dc_warp + cc < p.N) { -+ data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]); -+ } -+#endif // MUL_MAT_ID -+ } -+ } -+ } -+ } -+#endif // COOPMAT -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp -new file mode 100644 -index 00000000..cbfa5dce ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp -@@ -0,0 +1,328 @@ -+#version 450 -+ -+#extension GL_EXT_control_flow_attributes : enable -+#extension GL_EXT_shader_16bit_storage : require -+ -+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -+#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require -+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require -+ -+#extension GL_KHR_memory_scope_semantics : enable -+#extension GL_KHR_cooperative_matrix : enable -+#extension GL_NV_cooperative_matrix2 : enable -+#extension GL_EXT_buffer_reference : enable -+#extension GL_KHR_shader_subgroup_ballot : enable -+#extension GL_KHR_shader_subgroup_vote : enable -+ -+#include "types.comp" -+ -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+layout (constant_id = 1) const uint BM = 64; -+layout (constant_id = 2) const uint BN = 64; -+layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working with a quant -+ -+layout (push_constant) uniform parameter -+{ -+ uint M; -+ uint N; -+ uint K; -+ uint stride_a; -+ uint stride_b; -+ uint stride_d; -+ -+ uint batch_stride_a; -+ uint batch_stride_b; -+ uint batch_stride_d; -+ -+#ifdef MUL_MAT_ID -+ uint nei0; -+ uint nei1; -+ uint nbi1; -+ uint ne11; -+#else -+ uint k_split; -+ uint ne02; -+ uint ne12; -+ uint broadcast2; -+ uint broadcast3; -+#endif -+} p; -+ -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; -+layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; -+ -+#if QUANT_K > 1 -+#define DECODEFUNCA , dequantFuncA -+#define MAT_A_TYPE float16_t -+ -+#include "dequant_funcs_cm2.comp" -+ -+#else -+#define DECODEFUNCA -+#define MAT_A_TYPE A_TYPE -+#endif -+ -+#define MAT_B_TYPE B_TYPE -+ -+#ifdef MUL_MAT_ID -+layout (binding = 3) readonly buffer IDS {int data_ids[];}; -+ -+shared u16vec4 row_ids[3072]; -+ -+layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufB { -+ B_TYPE b[]; -+}; -+ -+uint _ne1; -+shared uint _ne1_sh; -+ -+B_TYPE decodeFuncB(const in decodeBufB bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -+{ -+ const uint row_i = blockCoords[0]; -+ -+ if (row_i >= _ne1) { -+ return B_TYPE(0.0); -+ } -+ -+ const u16vec4 row_idx = row_ids[row_i]; -+ B_TYPE ret = data_b[row_idx.y * p.batch_stride_b + row_idx.x * p.stride_b + blockCoords[1]]; -+ -+ return ret; -+} -+ -+D_TYPE perElemOpD(const in uint32_t r, const in uint32_t c, const in D_TYPE elem, const in uint32_t ir, const in uint32_t ic) -+{ -+ uint dr = ir * BM + r; -+ uint dc = ic * BN + c; -+ -+ if (dr < p.M && dc < _ne1) { -+ uint row_i = dc; -+ const u16vec4 row_idx = row_ids[row_i]; -+ data_d[row_idx.y * p.batch_stride_d + row_idx.z * p.stride_d + dr] = elem; -+ } -+ return elem; -+} -+ -+#endif -+ -+void main() { -+#if defined(DATA_A_IQ4_NL) -+ init_iq4nl_shmem(); -+#endif -+ -+#ifdef MUL_MAT_ID -+ const uint expert_idx = gl_GlobalInvocationID.z; -+#else -+ const uint batch_idx = gl_GlobalInvocationID.z; -+ -+ const uint i13 = batch_idx / p.ne12; -+ const uint i12 = batch_idx % p.ne12; -+ -+ const uint i03 = i13 / p.broadcast3; -+ const uint i02 = i12 / p.broadcast2; -+ -+ const uint batch_idx_a = i03 * p.ne02 + i02; -+#endif -+ -+ const uint blocks_m = (p.M + BM - 1) / BM; -+ const uint ir = gl_WorkGroupID.x % blocks_m; -+ const uint ik = gl_WorkGroupID.x / blocks_m; -+ const uint ic = gl_WorkGroupID.y; -+ -+#ifdef MUL_MAT_ID -+ // Spread the search across all elements in the first subgroup -+ if (gl_SubgroupID == 0) { -+ _ne1 = 0; -+ uint num_elements = p.nei1 * p.nei0; -+ -+ for (uint i = gl_SubgroupInvocationID; subgroupAny(i < num_elements); i += gl_SubgroupSize) { -+ bool in_range = i < num_elements; -+ uint ii0 = i % p.nei0; -+ uint ii1 = i / p.nei0; -+ uint id = in_range ? data_ids[ii1*p.nbi1 + ii0] : 0; -+ uvec4 ballot = subgroupBallot(in_range && id == expert_idx); -+ uint idx = subgroupBallotExclusiveBitCount(ballot); -+ if (in_range && id == expert_idx) { -+ row_ids[_ne1 + idx] = u16vec4(ii0 % p.ne11, ii1, ii0, 0); -+ } -+ _ne1 += subgroupBallotBitCount(ballot); -+ } -+ _ne1_sh = _ne1; -+ } -+ -+ barrier(); -+ -+ _ne1 = _ne1_sh; -+ -+ // Workgroup has no work -+ if (ic * BN >= _ne1) return; -+#endif -+ -+#ifdef MUL_MAT_ID -+ uint start_k = 0; -+ const uint end_k = p.K; -+#else -+ uint start_k = ik * p.k_split; -+ const uint end_k = min(p.K, (ik + 1) * p.k_split); -+#endif -+ -+ coopmat sum; -+ sum = coopmat(0.0); -+ -+#ifdef MUL_MAT_ID -+ uint pos_a = (expert_idx * p.batch_stride_a) / QUANT_K; -+ uint pos_b = 0; -+#else -+ uint pos_a = (batch_idx_a * p.batch_stride_a) / QUANT_K; -+ uint pos_b = batch_idx * p.batch_stride_b; -+#endif -+ -+ uint stride_a = p.stride_a / QUANT_K; -+ uint stride_b = p.stride_b; -+ -+ // Hint to the compiler that values are aligned (want 16B alignment). -+ // Quants are always block-aligned, no alignment needed. -+#if ALIGNED -+#if QUANT_K == 1 -+ stride_a &= ~7; -+#endif -+ stride_b &= ~7; -+#endif -+ -+ // Create layouts for both clamped and unclamped accesses -+ tensorLayoutNV<2> tensorLayoutA = createTensorLayoutNV(2); -+ tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutAClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); -+ tensorLayoutNV<2> tensorLayoutB = createTensorLayoutNV(2); -+ tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutBClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); -+ tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); -+ -+#if QUANT_K > 1 -+ tensorLayoutA = setTensorLayoutBlockSizeNV(tensorLayoutA, 1, QUANT_K); -+ tensorLayoutAClamp = setTensorLayoutBlockSizeNV(tensorLayoutAClamp, 1, QUANT_K); -+#endif -+ -+ // Use end_k rather than p.K as the dimension because that's what -+ // we need to bound check against when using split_k -+ tensorLayoutA = setTensorLayoutDimensionNV(tensorLayoutA, p.M, end_k); -+ tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, p.N, end_k); -+ tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.N, p.M); -+ tensorLayoutAClamp = setTensorLayoutDimensionNV(tensorLayoutAClamp, p.M, end_k); -+ tensorLayoutBClamp = setTensorLayoutDimensionNV(tensorLayoutBClamp, p.N, end_k); -+ -+ tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0); -+ -+#if !defined(MUL_MAT_ID) -+ // Detect a fast path where all loads are entirely in bounds and no clamping is required -+ if ((ir + 1) * BM <= p.M && (ic + 1) * BN <= p.N && (start_k % BK) == 0 && (end_k % BK) == 0 && -+#if QUANT_K == 1 -+ (stride_a % 8) == 0 && -+#endif -+ (stride_b % 8) == 0 && (start_k % 8) == 0) { -+ // Hint to the compiler that values are aligned (want 16B alignment) -+ start_k &= ~7; -+ stride_b &= ~7; -+#if QUANT_K == 1 -+ stride_a &= ~7; -+#endif -+ -+ tensorLayoutA = setTensorLayoutStrideNV(tensorLayoutA, stride_a, 1); -+ tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1); -+ -+ uint k_iters = (end_k - start_k + BK - 1) / BK; -+ -+ for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) { -+ -+ coopmat mat_a; -+ coopmat mat_b; -+ -+ coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA); -+ coopmat mat_a_ft = coopmat(mat_a); -+ -+ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose); -+ coopmat mat_b_ft = coopmat(mat_b); -+ -+ sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); -+ } -+ } else -+#endif // !defined(MUL_MAT_ID) -+ { -+ tensorLayoutA = setTensorLayoutStrideNV(tensorLayoutA, stride_a, 1); -+ -+ tensorLayoutAClamp = setTensorLayoutStrideNV(tensorLayoutAClamp, stride_a, 1); -+ -+ tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1); -+ -+ tensorLayoutBClamp = setTensorLayoutStrideNV(tensorLayoutBClamp, stride_b, 1); -+ -+ [[dont_unroll]] -+ for (uint block_k = start_k; block_k < end_k; block_k += BK) { -+ -+ coopmat mat_a; -+ coopmat mat_b; -+ coopmat mat_a_ft; -+ coopmat mat_b_ft; -+ -+ // Clamping is expensive, so detect different code paths for each combination -+ // of A and B needing clamping. -+ bool unclampedA = (ir + 1) * BM <= p.M && block_k + BK <= end_k && (block_k % 8) == 0; -+#ifdef MUL_MAT_ID -+ bool unclampedB = true; -+#else -+ bool unclampedB = (ic + 1) * BN <= p.N && block_k + BK <= end_k && (block_k % 8) == 0; -+#endif -+ if (unclampedA && unclampedB) { -+ coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, (block_k & ~7), BK) DECODEFUNCA); -+#ifdef MUL_MAT_ID -+ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose, decodeFuncB); -+#else -+ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, (block_k & ~7), BK), tensorViewTranspose); -+#endif -+ mat_a_ft = coopmat(mat_a); -+ mat_b_ft = coopmat(mat_b); -+ sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); -+ } else if (unclampedA && !unclampedB) { -+ coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, (block_k & ~7), BK) DECODEFUNCA); -+ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutBClamp, ic * BN, BN, block_k, BK), tensorViewTranspose); -+ -+ mat_a_ft = coopmat(mat_a); -+ mat_b_ft = coopmat(mat_b); -+ sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); -+ } else if (!unclampedA && unclampedB) { -+ coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutAClamp, ir * BM, BM, block_k, BK) DECODEFUNCA); -+#ifdef MUL_MAT_ID -+ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose, decodeFuncB); -+#else -+ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, (block_k & ~7), BK), tensorViewTranspose); -+#endif -+ mat_a_ft = coopmat(mat_a); -+ mat_b_ft = coopmat(mat_b); -+ sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); -+ } else if (!unclampedA && !unclampedB) { -+ coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutAClamp, ir * BM, BM, block_k, BK) DECODEFUNCA); -+ coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutBClamp, ic * BN, BN, block_k, BK), tensorViewTranspose); -+ -+ mat_a_ft = coopmat(mat_a); -+ mat_b_ft = coopmat(mat_b); -+ sum = coopMatMulAdd(mat_a_ft, mat_b_ft, sum); -+ } -+ } -+ } -+ -+ // Convert from ACC_TYPE to D_TYPE -+ coopmat mat_d; -+ mat_d = coopmat(sum); -+ -+#ifdef MUL_MAT_ID -+ // Call callback to store each element, remapping row through shared memory -+ coopMatPerElementNV(mat_d, mat_d, perElemOpD, ir, ic); -+#else -+ tensorLayoutD = setTensorLayoutStrideNV(tensorLayoutD, p.stride_d, 1); -+ -+ uint pos_d = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * gl_NumWorkGroups.z; -+ coopMatStoreTensorNV(mat_d, data_d, pos_d, sliceTensorLayoutNV(tensorLayoutD, ic * BN, BN, ir * BM, BM), tensorViewTranspose); -+#endif -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp -new file mode 100644 -index 00000000..6627a50b ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/norm.comp -@@ -0,0 +1,44 @@ -+#version 450 -+ -+#include "generic_head.comp" -+#include "types.comp" -+ -+#extension GL_EXT_control_flow_attributes : enable -+#define BLOCK_SIZE 512 -+ -+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+shared vec2 sum[BLOCK_SIZE]; -+ -+void main() { -+ const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; -+ const uint tid = gl_LocalInvocationID.x; -+ -+ sum[tid] = vec2(0.0f, 0.0f); -+ -+ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { -+ const float xi = float(data_a[row*p.KX + col]); -+ sum[tid].x += xi; -+ sum[tid].y += xi * xi; -+ } -+ -+ // sum up partial sums and write back result -+ barrier(); -+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { -+ if (tid < s) { -+ sum[tid] += sum[tid + s]; -+ } -+ barrier(); -+ } -+ -+ const float mean = sum[0].x / p.KX; -+ const float var = sum[0].y / p.KX - mean * mean; -+ const float inv_std = inversesqrt(var + p.param1); -+ -+ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { -+ data_d[row*p.KX + col] = D_TYPE((float(data_a[row*p.KX + col]) - mean) * inv_std); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp -new file mode 100644 -index 00000000..450b67fc ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pad.comp -@@ -0,0 +1,28 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_unary_head.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; -+ -+ if (idx >= p.ne) { -+ return; -+ } -+ -+ const uint i3 = idx / (p.ne12*p.ne11*p.ne10); -+ const uint i3_offset = i3 * p.ne12*p.ne11*p.ne10; -+ const uint i2 = (idx - i3_offset) / (p.ne11*p.ne10); -+ const uint i2_offset = i2*p.ne11*p.ne10; -+ const uint i1 = (idx - i3_offset - i2_offset) / p.ne10; -+ const uint i0 = idx - i3_offset - i2_offset - i1*p.ne10; -+ -+ const uint src0_idx = i3*p.nb03 + i2*p.nb02 + i1*p.nb01 + i0*p.nb00; -+ const uint dst_idx = i3*p.nb13 + i2*p.nb12 + i1*p.nb11 + i0*p.nb10; -+ -+ const bool is_src0 = i0 < p.ne00 && i1 < p.ne01 && i2 < p.ne02 && i3 < p.ne03; -+ -+ data_d[get_doffset() + dst_idx] = D_TYPE(is_src0 ? data_a[get_aoffset() + src0_idx] : 0.0f); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp -new file mode 100644 -index 00000000..b6124411 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp -@@ -0,0 +1,74 @@ -+#version 450 -+ -+#include "types.comp" -+ -+#extension GL_EXT_shader_16bit_storage : require -+ -+layout(push_constant) uniform parameter { -+ uint IW; uint IH; -+ uint OW; uint OH; -+ uint OC; -+ uint pelements; -+ uint op; -+ int k0; int k1; -+ int s0; int s1; -+ int p0; int p1; -+} p; -+ -+#define BLOCK_SIZE 512 -+#define FLT_MAX 3.402823466e+38F -+#define OP_POOL_MAX 0u -+#define OP_POOL_AVG 1u -+ -+layout (local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; -+ -+layout(binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout(binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+void main() { -+ const uint idx = gl_GlobalInvocationID.x; -+ if (idx >= p.pelements) { -+ return; -+ } -+ -+ const uint O_HW = p.OW * p.OH; -+ -+ const uint nc = idx / O_HW; -+ const uint cur_oh = (idx % O_HW) / p.OW; -+ const uint cur_ow = (idx % O_HW) % p.OW; -+ -+ const int start_h = int(cur_oh) * p.s0 - p.p0; -+ const uint bh = max(start_h, 0); -+ const uint eh = min(start_h + p.k0, p.IH); -+ -+ const int start_w = int(cur_ow) * p.s1 - p.p1; -+ const uint bw = max(start_w, 0); -+ const uint ew = min(start_w + p.k1, p.IW); -+ -+ const float scale = 1.0 / float(p.k0 * p.k1); -+ float res; -+ -+ if (p.op == OP_POOL_AVG) { -+ res = 0.0; -+ } else if (p.op == OP_POOL_MAX) { -+ res = -FLT_MAX; -+ } else { -+ return; -+ } -+ -+ #pragma unroll -+ for (uint i = bh; i < eh; i++) { -+ #pragma unroll -+ for (uint j = bw; j < ew; j++) { -+ const float cur = D_TYPE(data_a[nc * p.IH * p.IW + i * p.IW + j]); -+ -+ if (p.op == OP_POOL_AVG) { -+ res += cur * scale; -+ } else if (p.op == OP_POOL_MAX) { -+ res = max(res, cur); -+ } -+ } -+ } -+ -+ data_d[nc * O_HW + cur_oh * p.OW + cur_ow] = res; -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp -new file mode 100644 -index 00000000..52a19b62 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/relu.comp -@@ -0,0 +1,21 @@ -+#version 450 -+ -+#include "generic_head.comp" -+#include "types.comp" -+ -+#extension GL_EXT_control_flow_attributes : enable -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+void main() { -+ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; -+ -+ if (i >= p.KX) { -+ return; -+ } -+ -+ data_d[i] = max(float(data_a[i]), 0); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp -new file mode 100644 -index 00000000..1568b141 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/repeat.comp -@@ -0,0 +1,26 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_unary_head.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+uint src0_idx_mod(uint idx) { -+ const uint i13 = idx / (p.ne12*p.ne11*p.ne10); -+ const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10; -+ const uint i12 = (idx - i13_offset) / (p.ne11*p.ne10); -+ const uint i12_offset = i12*p.ne11*p.ne10; -+ const uint i11 = (idx - i13_offset - i12_offset) / p.ne10; -+ const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10; -+ return (i13 % p.ne03)*p.nb03 + (i12 % p.ne02)*p.nb02 + (i11 % p.ne01)*p.nb01 + (i10 % p.ne00)*p.nb00; -+} -+ -+void main() { -+ const uint idx = get_idx(); -+ -+ if (idx >= p.ne) { -+ return; -+ } -+ -+ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(data_a[get_aoffset() + src0_idx_mod(idx)]); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp -new file mode 100644 -index 00000000..b554400b ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp -@@ -0,0 +1,42 @@ -+#version 450 -+ -+#include "generic_head.comp" -+#include "types.comp" -+ -+#extension GL_EXT_control_flow_attributes : enable -+#define BLOCK_SIZE 512 -+ -+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+shared FLOAT_TYPE sum[BLOCK_SIZE]; -+ -+void main() { -+ const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; -+ const uint tid = gl_LocalInvocationID.x; -+ -+ sum[tid] = FLOAT_TYPE(0.0f); // partial sum for thread in warp -+ -+ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { -+ const FLOAT_TYPE xi = FLOAT_TYPE(data_a[row*p.KX + col]); -+ sum[tid] += xi * xi; -+ } -+ -+ // sum up partial sums and write back result -+ barrier(); -+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { -+ if (tid < s) { -+ sum[tid] += sum[tid + s]; -+ } -+ barrier(); -+ } -+ -+ const FLOAT_TYPE mean = sum[0] / FLOAT_TYPE(p.KX); -+ const FLOAT_TYPE scale = inversesqrt(mean + FLOAT_TYPE(p.param1)); -+ -+ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) { -+ data_d[row*p.KX + col] = D_TYPE(scale * FLOAT_TYPE(data_a[row*p.KX + col])); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp -new file mode 100644 -index 00000000..574b51ca ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_head.comp -@@ -0,0 +1,49 @@ -+#include "types.comp" -+ -+#extension GL_EXT_shader_16bit_storage : require -+#extension GL_EXT_spirv_intrinsics: enable -+ -+#if RTE16 -+spirv_execution_mode(capabilities = [4467], 4462, 16); // RoundingModeRTE, 16 bits -+#endif -+ -+layout(local_size_x = 1, local_size_y = 256, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) readonly buffer Y {int data_pos[];}; -+layout (binding = 2) readonly buffer Z {float data_ff[];}; -+layout (binding = 3) writeonly buffer D {D_TYPE data_d[];}; -+ -+layout (push_constant) uniform parameter { -+ uint ncols; -+ uint n_dims; -+ float freq_scale; -+ uint p_delta_rows; -+ float freq_base; -+ float ext_factor; -+ float attn_factor; -+ float corr_dims[2]; -+ float theta_scale; -+ uint has_ff; -+} p; -+ -+float rope_yarn_ramp(const float low, const float high, const uint i0) { -+ const float y = (i0 / 2 - low) / max(0.001f, high - low); -+ return 1.0f - min(1.0f, max(0.0f, y)); -+} -+ -+void rope_yarn(const float theta_extrap, const uint i0, out float cos_theta, out float sin_theta) { -+ float mscale = p.attn_factor; -+ // Get n-d rotational scaling corrected for extrapolation -+ float theta_interp = p.freq_scale * theta_extrap; -+ float theta = theta_interp; -+ if (p.ext_factor != 0.0f) { -+ float ramp_mix = rope_yarn_ramp(p.corr_dims[0], p.corr_dims[1], i0) * p.ext_factor; -+ theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix; -+ -+ // Get n-d magnitude scaling corrected for interpolation -+ mscale *= 1.0f + 0.1f * log(1.0f / p.freq_scale); -+ } -+ cos_theta = cos(theta) * mscale; -+ sin_theta = sin(theta) * mscale; -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp -new file mode 100644 -index 00000000..83b46b69 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp -@@ -0,0 +1,37 @@ -+#version 450 -+ -+#include "rope_head.comp" -+ -+void main() { -+ const uint col = gl_GlobalInvocationID.y * 2; -+ const uint row = gl_GlobalInvocationID.x; -+ -+ if (col >= p.ncols) { -+ return; -+ } -+ -+ if (col >= p.n_dims) { -+ const uint i = row*p.ncols + col; -+ -+ data_d[i + 0] = data_a[i + 0]; -+ data_d[i + 1] = data_a[i + 1]; -+ -+ return; -+ } -+ -+ const uint i = row*p.ncols + col/2; -+ const uint i2 = row/p.p_delta_rows; -+ -+ const float theta_base = data_pos[i2] * pow(p.theta_scale, col/2.0f); -+ -+ const float freq_factor = p.has_ff != 0 ? data_ff[col/2] : 1.0f; -+ -+ float cos_theta, sin_theta; -+ rope_yarn(theta_base / freq_factor, col, cos_theta, sin_theta); -+ -+ const float x0 = float(data_a[i + 0]); -+ const float x1 = float(data_a[i + p.n_dims/2]); -+ -+ data_d[i + 0] = D_TYPE(x0*cos_theta - x1*sin_theta); -+ data_d[i + p.n_dims/2] = D_TYPE(x0*sin_theta + x1*cos_theta); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp -new file mode 100644 -index 00000000..e416ad93 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp -@@ -0,0 +1,37 @@ -+#version 450 -+ -+#include "rope_head.comp" -+ -+void main() { -+ const uint col = gl_GlobalInvocationID.y * 2; -+ const uint row = gl_GlobalInvocationID.x; -+ -+ if (col >= p.ncols) { -+ return; -+ } -+ -+ if (col >= p.n_dims) { -+ const uint i = row*p.ncols + col; -+ -+ data_d[i + 0] = data_a[i + 0]; -+ data_d[i + 1] = data_a[i + 1]; -+ -+ return; -+ } -+ -+ const uint i = row*p.ncols + col; -+ const uint i2 = row/p.p_delta_rows; -+ -+ const float theta_base = data_pos[i2] * pow(p.theta_scale, col/2.0f); -+ -+ const float freq_factor = p.has_ff != 0 ? data_ff[col/2] : 1.0f; -+ -+ float cos_theta, sin_theta; -+ rope_yarn(theta_base / freq_factor, col, cos_theta, sin_theta); -+ -+ const float x0 = float(data_a[i + 0]); -+ const float x1 = float(data_a[i + 1]); -+ -+ data_d[i + 0] = D_TYPE(x0*cos_theta - x1*sin_theta); -+ data_d[i + 1] = D_TYPE(x0*sin_theta + x1*cos_theta); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp -new file mode 100644 -index 00000000..4663428d ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/scale.comp -@@ -0,0 +1,24 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_unary_head.comp" -+ -+const uint num_threads = 128; -+ -+layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ uint idx = get_idx(); -+ -+ // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation -+ const uint num_iter = 4; -+ -+ [[unroll]] for (uint i = 0; i < num_iter; ++i) { -+ if (idx >= p.ne) { -+ continue; -+ } -+ -+ data_d[get_doffset() + idx] = D_TYPE(FLOAT_TYPE(data_a[get_aoffset() + idx]) * FLOAT_TYPE(p.param1)); -+ idx += num_threads; -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp -new file mode 100644 -index 00000000..4d36f88e ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/silu.comp -@@ -0,0 +1,22 @@ -+#version 450 -+ -+#include "generic_head.comp" -+#include "types.comp" -+ -+#extension GL_EXT_control_flow_attributes : enable -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+void main() { -+ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; -+ -+ if (i >= p.KX) { -+ return; -+ } -+ -+ const float xi = float(data_a[i]); -+ data_d[i] = D_TYPE(xi / (1.0f + exp(-xi))); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp -new file mode 100644 -index 00000000..d7c15a16 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sin.comp -@@ -0,0 +1,17 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_unary_head.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ const uint idx = get_idx(); -+ -+ if (idx >= p.ne) { -+ return; -+ } -+ -+ const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); -+ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(sin(val)); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp -new file mode 100644 -index 00000000..a25808e1 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp -@@ -0,0 +1,174 @@ -+#version 450 -+ -+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -+#extension GL_EXT_control_flow_attributes : enable -+ -+layout (push_constant) uniform parameter -+{ -+ uint KX; -+ uint KY; -+ float scale; -+ float max_bias; -+ float m0; -+ float m1; -+ uint n_head_log2; -+ uint nrows_x; -+} p; -+ -+#include "types.comp" -+ -+layout(constant_id = 0) const uint BLOCK_SIZE = 32; -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) readonly buffer Y {B_TYPE data_b[];}; -+layout (binding = 2) buffer D {D_TYPE data_d[];}; -+ -+shared FLOAT_TYPE vals[BLOCK_SIZE]; -+ -+// num_iters is the number of BLOCK_SIZE loop iterations we need to iterate -+// over all the columns. The main function tries to pass a constant here, -+// as if it were a template function, to allow unrolling. -+void soft_max(uint num_iters) { -+ const uint tid = gl_LocalInvocationID.x; -+ const uint rowx = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; -+ const uint rowy = (p.KY > 0) ? (rowx % p.KY) : 0; -+ -+ if (rowx >= p.nrows_x) { -+ return; -+ } -+ -+ float slope = 1.0f; -+ -+ // ALiBi -+ if (p.max_bias > 0.0f) { -+ const uint h = rowx/p.KY; // head index -+ -+ const float base = h < p.n_head_log2 ? p.m0 : p.m1; -+ const uint exp = h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1; -+ -+ slope = pow(base, exp); -+ } -+ -+ // Find max -+ FLOAT_TYPE max_val = uintBitsToFloat(0xFF800000); -+ -+ // Cache values while we compute the max, so we don't need to read them -+ // again when we're ready to compute exp(x-max). -+ const uint DATA_CACHE_SIZE = 16; -+ FLOAT_TYPE data_cache[DATA_CACHE_SIZE]; -+ -+ [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { -+ const uint col = col0 + tid; -+ -+ FLOAT_TYPE a = FLOAT_TYPE(0); -+ if (col < p.KX) { -+ a = data_a[rowx * p.KX + col]; -+ } -+ -+ FLOAT_TYPE b = FLOAT_TYPE(0); -+ if (p.KY > 0 && col < p.KX) { -+ b = data_b[rowy * p.KX + col]; -+ } -+ -+ FLOAT_TYPE v = a * p.scale + slope * b; -+ -+ if (col < p.KX) { -+ max_val = max(max_val, v); -+ } -+ -+ if (idx < DATA_CACHE_SIZE) { -+ data_cache[idx] = v; -+ } -+ } -+ -+ // reduce across the workgroup -+ vals[tid] = max_val; -+ barrier(); -+ [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { -+ if (tid < s) { -+ vals[tid] = max(vals[tid], vals[tid + s]); -+ } -+ barrier(); -+ } -+ -+ max_val = vals[0]; -+ barrier(); -+ -+ FLOAT_TYPE sum = FLOAT_TYPE(0.0f); -+ -+ // Compute sum{exp(x - max)} -+ [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { -+ const uint col = col0 + tid; -+ -+ if (col >= p.KX) { -+ break; -+ } -+ -+ // compute exp(a*scale+b*slope), add it to sum, and cache the new value -+ // in data_cache if possible. -+ const uint i = rowx * p.KX + col; -+ FLOAT_TYPE val; -+ if (idx < DATA_CACHE_SIZE) { -+ val = exp(data_cache[idx] - max_val); -+ } else { -+ val = exp(FLOAT_TYPE(data_a[i]) * p.scale + (p.KY > 0 ? slope * FLOAT_TYPE(data_b[rowy * p.KX + col]) : FLOAT_TYPE(0.0f)) - max_val); -+ } -+ sum += val; -+ if (idx < DATA_CACHE_SIZE) { -+ data_cache[idx] = val; -+ } else { -+ data_d[i] = D_TYPE(val); -+ } -+ } -+ -+ // reduce across the workgroup -+ vals[tid] = sum; -+ barrier(); -+ [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { -+ if (tid < s) { -+ vals[tid] += vals[tid + s]; -+ } -+ barrier(); -+ } -+ sum = vals[0]; -+ -+ FLOAT_TYPE rcpdivisor = 1.0/sum; -+ -+ [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) { -+ const uint col = col0 + tid; -+ -+ if (col >= p.KX) { -+ continue; -+ } -+ -+ if (idx < DATA_CACHE_SIZE) { -+ data_d[rowx*p.KX + col] = D_TYPE(data_cache[idx] * rcpdivisor); -+ } else { -+ data_d[rowx*p.KX + col] *= D_TYPE(rcpdivisor); -+ } -+ } -+} -+ -+void main() { -+ // instantiate the soft_max function for several different -+ // dimensions, to allow loop unrolling -+ uint num_blocks = (p.KX + BLOCK_SIZE - 1) / BLOCK_SIZE; -+ if (num_blocks > 32) { -+ soft_max(num_blocks); -+ } else if (num_blocks > 16) { -+ soft_max(32); -+ } else if (num_blocks > 8) { -+ soft_max(16); -+ } else if (num_blocks > 4) { -+ soft_max(8); -+ } else if (num_blocks == 4) { -+ soft_max(4); -+ } else if (num_blocks == 3) { -+ soft_max(3); -+ } else if (num_blocks == 2) { -+ soft_max(2); -+ } else if (num_blocks == 1) { -+ soft_max(1); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp -new file mode 100644 -index 00000000..ef43598b ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/square.comp -@@ -0,0 +1,17 @@ -+#version 450 -+ -+#include "types.comp" -+#include "generic_unary_head.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+void main() { -+ const uint idx = get_idx(); -+ -+ if (idx >= p.ne) { -+ return; -+ } -+ -+ const FLOAT_TYPE val = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(idx)]); -+ data_d[get_doffset() + dst_idx(idx)] = D_TYPE(val * val); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp -new file mode 100644 -index 00000000..961e5ffa ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.comp -@@ -0,0 +1,37 @@ -+#version 450 -+ -+#include "generic_head.comp" -+#include "types.comp" -+ -+#extension GL_EXT_control_flow_attributes : enable -+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+layout (constant_id = 0) const uint BLOCK_SIZE = 32; -+ -+shared FLOAT_TYPE tmp[BLOCK_SIZE]; -+ -+void main() { -+ const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; -+ const uint col = gl_LocalInvocationID.x; -+ -+ tmp[col] = FLOAT_TYPE(0.0f); -+ -+ for (uint i = col; i < p.KX; i += BLOCK_SIZE) { -+ tmp[col] += FLOAT_TYPE(data_a[row*p.KX + i]); -+ } -+ -+ barrier(); -+ [[unroll]] for (int s = int(BLOCK_SIZE) / 2; s > 0; s >>= 1) { -+ if (col < s) { -+ tmp[col] += tmp[col + s]; -+ } -+ barrier(); -+ } -+ -+ if (col == 0) { -+ data_d[row] = D_TYPE(tmp[0]); -+ } -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp -new file mode 100644 -index 00000000..495f966b ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/tanh.comp -@@ -0,0 +1,20 @@ -+#version 450 -+ -+#include "generic_head.comp" -+#include "types.comp" -+ -+#extension GL_EXT_control_flow_attributes : enable -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+void main() { -+ const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; -+ -+ if (i >= p.KX) { -+ return; -+ } -+ data_d[i] = D_TYPE(1. - 2. / (exp(2.*data_a[i]) + 1.)); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp -new file mode 100644 -index 00000000..28eb24e1 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/test_coopmat2_support.comp -@@ -0,0 +1,7 @@ -+#version 460 -+ -+#extension GL_NV_cooperative_matrix2 : require -+ -+void main() -+{ -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp -new file mode 100644 -index 00000000..79e065a9 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/timestep_embedding.comp -@@ -0,0 +1,41 @@ -+#version 450 -+ -+#extension GL_EXT_shader_16bit_storage : require -+ -+layout (push_constant) uniform parameter -+{ -+ uint nb1; -+ uint dim; -+ uint max_period; -+} p; -+ -+#include "types.comp" -+ -+#extension GL_EXT_control_flow_attributes : enable -+#define BLOCK_SIZE 256 -+ -+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+void main() { -+ const uint i = gl_WorkGroupID.y; -+ const uint j = gl_GlobalInvocationID.x; -+ const uint d_offset = i * p.nb1; -+ -+ if (p.dim % 2 != 0 && j == ((p.dim + 1) / 2)) { -+ data_d[d_offset + p.dim] = 0.f; -+ } -+ -+ const uint half_dim = p.dim / 2; -+ if (j >= half_dim) { -+ return; -+ } -+ -+ const float timestep = float(data_a[i]); -+ const float freq = float(exp(-log(p.max_period) * j / half_dim)); -+ const float arg = timestep * freq; -+ data_d[d_offset + j] = D_TYPE(cos(arg)); -+ data_d[d_offset + j + half_dim] = D_TYPE(sin(arg)); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp -new file mode 100644 -index 00000000..eecc47f3 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/types.comp -@@ -0,0 +1,323 @@ -+ -+#if !defined(GGML_TYPES_COMP) -+#define GGML_TYPES_COMP -+ -+#extension GL_EXT_shader_explicit_arithmetic_types : require -+ -+#if defined(DATA_A_F32) -+#define QUANT_K 1 -+#define QUANT_R 1 -+ -+#if !defined(LOAD_VEC_A) || LOAD_VEC_A == 1 -+#define A_TYPE float -+#elif LOAD_VEC_A == 4 -+#define A_TYPE vec4 -+#elif LOAD_VEC_A == 8 -+#define A_TYPE mat2x4 -+#endif -+#endif -+ -+#if defined(DATA_A_F16) -+#define QUANT_K 1 -+#define QUANT_R 1 -+ -+#if !defined(LOAD_VEC_A) || LOAD_VEC_A == 1 -+#define A_TYPE float16_t -+#elif LOAD_VEC_A == 4 -+#define A_TYPE f16vec4 -+#elif LOAD_VEC_A == 8 -+#define A_TYPE f16mat2x4 -+#endif -+#endif -+ -+#define QUANT_K_Q4_0 32 -+#define QUANT_R_Q4_0 2 -+ -+struct block_q4_0 -+{ -+ float16_t d; -+ uint8_t qs[16]; -+}; -+struct block_q4_0_packed16 -+{ -+ float16_t d; -+ uint16_t qs[16/2]; -+}; -+ -+#if defined(DATA_A_Q4_0) -+#define QUANT_K QUANT_K_Q4_0 -+#define QUANT_R QUANT_R_Q4_0 -+#define A_TYPE block_q4_0 -+#define A_TYPE_PACKED16 block_q4_0_packed16 -+#endif -+ -+#define QUANT_K_Q4_1 32 -+#define QUANT_R_Q4_1 2 -+ -+struct block_q4_1 -+{ -+ float16_t d; -+ float16_t m; -+ uint8_t qs[16]; -+}; -+ -+struct block_q4_1_packed16 -+{ -+ float16_t d; -+ float16_t m; -+ uint16_t qs[16/2]; -+}; -+ -+#if defined(DATA_A_Q4_1) -+#define QUANT_K QUANT_K_Q4_1 -+#define QUANT_R QUANT_R_Q4_1 -+#define A_TYPE block_q4_1 -+#define A_TYPE_PACKED16 block_q4_1_packed16 -+#endif -+ -+#define QUANT_K_Q5_0 32 -+#define QUANT_R_Q5_0 2 -+ -+struct block_q5_0 -+{ -+ float16_t d; -+ uint16_t qh[2]; -+ uint8_t qs[16]; -+}; -+ -+struct block_q5_0_packed16 -+{ -+ float16_t d; -+ uint16_t qh[2]; -+ uint16_t qs[16/2]; -+}; -+ -+#if defined(DATA_A_Q5_0) -+#define QUANT_K QUANT_K_Q5_0 -+#define QUANT_R QUANT_R_Q5_0 -+#define A_TYPE block_q5_0 -+#define A_TYPE_PACKED16 block_q5_0_packed16 -+#endif -+ -+#define QUANT_K_Q5_1 32 -+#define QUANT_R_Q5_1 2 -+ -+struct block_q5_1 -+{ -+ float16_t d; -+ float16_t m; -+ uint qh; -+ uint8_t qs[16]; -+}; -+ -+struct block_q5_1_packed16 -+{ -+ float16_t d; -+ float16_t m; -+ uint qh; -+ uint16_t qs[16/2]; -+}; -+ -+#if defined(DATA_A_Q5_1) -+#define QUANT_K QUANT_K_Q5_1 -+#define QUANT_R QUANT_R_Q5_1 -+#define A_TYPE block_q5_1 -+#define A_TYPE_PACKED16 block_q5_1_packed16 -+#endif -+ -+#define QUANT_K_Q8_0 32 -+#define QUANT_R_Q8_0 1 -+ -+struct block_q8_0 -+{ -+ float16_t d; -+ int8_t qs[32]; -+}; -+struct block_q8_0_packed16 -+{ -+ float16_t d; -+ uint16_t qs[32/2]; -+}; -+ -+#if defined(DATA_A_Q8_0) -+#define QUANT_K QUANT_K_Q8_0 -+#define QUANT_R QUANT_R_Q8_0 -+#define A_TYPE block_q8_0 -+#define A_TYPE_PACKED16 block_q8_0_packed16 -+#endif -+ -+// K-quants -+#define QUANT_K_Q2_K 256 -+ -+struct block_q2_K -+{ -+ uint8_t scales[QUANT_K_Q2_K/16]; -+ uint8_t qs[QUANT_K_Q2_K/4]; -+ f16vec2 d; -+}; -+ -+struct block_q2_K_packed16 -+{ -+ uint16_t scales[QUANT_K_Q2_K/16/2]; -+ uint16_t qs[QUANT_K_Q2_K/4/2]; -+ f16vec2 d; -+}; -+ -+struct block_q2_K_packed32 -+{ -+ uint32_t scales[QUANT_K_Q2_K/16/4]; -+ uint32_t qs[QUANT_K_Q2_K/4/4]; -+ f16vec2 d; -+}; -+ -+#if defined(DATA_A_Q2_K) -+#define QUANT_K QUANT_K_Q2_K -+#define A_TYPE block_q2_K -+#define A_TYPE_PACKED16 block_q2_K_packed16 -+#define A_TYPE_PACKED32 block_q2_K_packed32 -+#endif -+ -+#define QUANT_K_Q3_K 256 -+ -+struct block_q3_K -+{ -+ uint8_t hmask[QUANT_K_Q3_K/8]; -+ uint8_t qs[QUANT_K_Q3_K/4]; -+ uint8_t scales[12]; -+ float16_t d; -+}; -+ -+struct block_q3_K_packed16 -+{ -+ uint16_t hmask[QUANT_K_Q3_K/8/2]; -+ uint16_t qs[QUANT_K_Q3_K/4/2]; -+ uint16_t scales[12/2]; -+ float16_t d; -+}; -+ -+#if defined(DATA_A_Q3_K) -+#define QUANT_K QUANT_K_Q3_K -+#define A_TYPE block_q3_K -+#define A_TYPE_PACKED16 block_q3_K_packed16 -+#endif -+ -+#define QUANT_K_Q4_K 256 -+ -+struct block_q4_K -+{ -+ f16vec2 d; -+ uint8_t scales[3*QUANT_K_Q4_K/64]; -+ uint8_t qs[QUANT_K_Q4_K/2]; -+}; -+ -+struct block_q4_K_packed16 -+{ -+ f16vec2 d; -+ uint16_t scales[3*QUANT_K_Q4_K/64/2]; -+ uint16_t qs[QUANT_K_Q4_K/2/2]; -+}; -+ -+struct block_q4_K_packed32 -+{ -+ f16vec2 d; -+ uint32_t scales[3*QUANT_K_Q4_K/64/4]; -+ uint32_t qs[QUANT_K_Q4_K/2/4]; -+}; -+ -+#if defined(DATA_A_Q4_K) -+#define QUANT_K QUANT_K_Q4_K -+#define A_TYPE block_q4_K -+#define A_TYPE_PACKED16 block_q4_K_packed16 -+#define A_TYPE_PACKED32 block_q4_K_packed32 -+#endif -+ -+#define QUANT_K_Q5_K 256 -+ -+struct block_q5_K -+{ -+ f16vec2 d; -+ uint8_t scales[12]; -+ uint8_t qh[QUANT_K_Q5_K/8]; -+ uint8_t qs[QUANT_K_Q5_K/2]; -+}; -+ -+struct block_q5_K_packed16 -+{ -+ f16vec2 d; -+ uint16_t scales[12/2]; -+ uint16_t qh[QUANT_K_Q5_K/8/2]; -+ uint16_t qs[QUANT_K_Q5_K/2/2]; -+}; -+ -+#if defined(DATA_A_Q5_K) -+#define QUANT_K QUANT_K_Q5_K -+#define A_TYPE block_q5_K -+#define A_TYPE_PACKED16 block_q5_K_packed16 -+#endif -+ -+#define QUANT_K_Q6_K 256 -+ -+struct block_q6_K -+{ -+ uint8_t ql[QUANT_K_Q6_K/2]; -+ uint8_t qh[QUANT_K_Q6_K/4]; -+ int8_t scales[QUANT_K_Q6_K/16]; -+ float16_t d; -+}; -+ -+struct block_q6_K_packed16 -+{ -+ uint16_t ql[QUANT_K_Q6_K/2/2]; -+ uint16_t qh[QUANT_K_Q6_K/4/2]; -+ int8_t scales[QUANT_K_Q6_K/16]; -+ float16_t d; -+}; -+ -+#if defined(DATA_A_Q6_K) -+#define QUANT_K QUANT_K_Q6_K -+#define A_TYPE block_q6_K -+#define A_TYPE_PACKED16 block_q6_K_packed16 -+#endif -+ -+// IQuants -+ -+#define QUANT_K_IQ4_NL 32 -+#define QUANT_R_IQ4_NL 2 -+ -+struct block_iq4_nl -+{ -+ float16_t d; -+ uint8_t qs[QUANT_K_IQ4_NL/2]; -+}; -+ -+struct block_iq4_nl_packed16 -+{ -+ float16_t d; -+ uint16_t qs[QUANT_K_IQ4_NL/2/2]; -+}; -+ -+#if defined(DATA_A_IQ4_NL) -+ -+const int8_t kvalues_iq4nl_const[16] = { -+ int8_t(-127), int8_t(-104), int8_t(-83), int8_t(-65), int8_t(-49), int8_t(-35), int8_t(-22), int8_t(-10), -+ int8_t(1), int8_t(13), int8_t(25), int8_t(38), int8_t(53), int8_t(69), int8_t(89), int8_t(113) -+}; -+ -+shared FLOAT_TYPE kvalues_iq4nl[16]; -+ -+void init_iq4nl_shmem() -+{ -+ // copy the table into shared memory and sync -+ if (gl_LocalInvocationIndex.x < 16) { -+ kvalues_iq4nl[gl_LocalInvocationIndex.x] = FLOAT_TYPE(kvalues_iq4nl_const[gl_LocalInvocationIndex.x]); -+ } -+ barrier(); -+} -+ -+#define QUANT_K QUANT_K_IQ4_NL -+#define QUANT_R QUANT_R_IQ4_NL -+#define A_TYPE block_iq4_nl -+#define A_TYPE_PACKED16 block_iq4_nl_packed16 -+#endif -+ -+#endif // !defined(GGML_TYPES_COMP) -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp -new file mode 100644 -index 00000000..6f607380 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp -@@ -0,0 +1,36 @@ -+#version 450 -+ -+layout (push_constant) uniform parameter -+{ -+ uint ne; uint a_offset; uint d_offset; -+ uint nb00; uint nb01; uint nb02; uint nb03; -+ uint ne10; uint ne11; uint ne12; uint ne13; -+ float sf0; float sf1; float sf2; float sf3; -+} p; -+ -+#include "types.comp" -+ -+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -+ -+layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -+ -+void main() { -+ const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; -+ -+ if (idx >= p.ne) { -+ return; -+ } -+ -+ const uint i10 = idx % p.ne10; -+ const uint i11 = (idx / p.ne10) % p.ne11; -+ const uint i12 = (idx / (p.ne10 * p.ne11)) % p.ne12; -+ const uint i13 = (idx / (p.ne10 * p.ne11 * p.ne12)) % p.ne13; -+ -+ const uint i00 = uint(i10 / p.sf0); -+ const uint i01 = uint(i11 / p.sf1); -+ const uint i02 = uint(i12 / p.sf2); -+ const uint i03 = uint(i13 / p.sf3); -+ -+ data_d[p.d_offset + idx] = D_TYPE(data_a[p.a_offset + i03 * p.nb03 + i02 * p.nb02 + i01 * p.nb01 + i00 * p.nb00]); -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp -new file mode 100644 -index 00000000..8111c063 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp -@@ -0,0 +1,594 @@ -+ -+ -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+ -+#ifdef _WIN32 -+ #include -+ #include // For _mkdir on Windows -+ #include // For std::replace on w64devkit -+#else -+ #include -+ #include -+ #include -+#endif -+ -+#include -+ -+#define ASYNCIO_CONCURRENCY 64 -+ -+std::mutex lock; -+std::vector> shader_fnames; -+ -+std::string GLSLC = "glslc"; -+std::string input_dir = "vulkan-shaders"; -+std::string output_dir = "/tmp"; -+std::string target_hpp = "ggml-vulkan-shaders.hpp"; -+std::string target_cpp = "ggml-vulkan-shaders.cpp"; -+bool no_clean = false; -+ -+const std::vector type_names = { -+ "f32", -+ "f16", -+ "q4_0", -+ "q4_1", -+ "q5_0", -+ "q5_1", -+ "q8_0", -+ "q2_k", -+ "q3_k", -+ "q4_k", -+ "q5_k", -+ "q6_k", -+ "iq4_nl" -+}; -+ -+namespace { -+void execute_command(const std::string& command, std::string& stdout_str, std::string& stderr_str) { -+#ifdef _WIN32 -+ HANDLE stdout_read, stdout_write; -+ HANDLE stderr_read, stderr_write; -+ SECURITY_ATTRIBUTES sa = { sizeof(SECURITY_ATTRIBUTES), NULL, TRUE }; -+ -+ if (!CreatePipe(&stdout_read, &stdout_write, &sa, 0) || -+ !SetHandleInformation(stdout_read, HANDLE_FLAG_INHERIT, 0)) { -+ throw std::runtime_error("Failed to create stdout pipe"); -+ } -+ -+ if (!CreatePipe(&stderr_read, &stderr_write, &sa, 0) || -+ !SetHandleInformation(stderr_read, HANDLE_FLAG_INHERIT, 0)) { -+ throw std::runtime_error("Failed to create stderr pipe"); -+ } -+ -+ PROCESS_INFORMATION pi; -+ STARTUPINFOA si = {}; -+ si.cb = sizeof(STARTUPINFOA); -+ si.dwFlags = STARTF_USESTDHANDLES; -+ si.hStdOutput = stdout_write; -+ si.hStdError = stderr_write; -+ -+ std::vector cmd(command.begin(), command.end()); -+ cmd.push_back('\0'); -+ -+ if (!CreateProcessA(NULL, cmd.data(), NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi)) { -+ throw std::runtime_error("Failed to create process"); -+ } -+ -+ CloseHandle(stdout_write); -+ CloseHandle(stderr_write); -+ -+ std::array buffer; -+ DWORD bytes_read; -+ -+ while (ReadFile(stdout_read, buffer.data(), (DWORD)buffer.size(), &bytes_read, NULL) && bytes_read > 0) { -+ stdout_str.append(buffer.data(), bytes_read); -+ } -+ -+ while (ReadFile(stderr_read, buffer.data(), (DWORD)buffer.size(), &bytes_read, NULL) && bytes_read > 0) { -+ stderr_str.append(buffer.data(), bytes_read); -+ } -+ -+ CloseHandle(stdout_read); -+ CloseHandle(stderr_read); -+ WaitForSingleObject(pi.hProcess, INFINITE); -+ CloseHandle(pi.hProcess); -+ CloseHandle(pi.hThread); -+#else -+int stdout_pipe[2]; -+ int stderr_pipe[2]; -+ -+ if (pipe(stdout_pipe) != 0 || pipe(stderr_pipe) != 0) { -+ throw std::runtime_error("Failed to create pipes"); -+ } -+ -+ pid_t pid = fork(); -+ if (pid < 0) { -+ throw std::runtime_error("Failed to fork process"); -+ } -+ -+ if (pid == 0) { -+ close(stdout_pipe[0]); -+ close(stderr_pipe[0]); -+ dup2(stdout_pipe[1], STDOUT_FILENO); -+ dup2(stderr_pipe[1], STDERR_FILENO); -+ close(stdout_pipe[1]); -+ close(stderr_pipe[1]); -+ execl("/bin/sh", "sh", "-c", command.c_str(), (char*) nullptr); -+ _exit(EXIT_FAILURE); -+ } else { -+ close(stdout_pipe[1]); -+ close(stderr_pipe[1]); -+ -+ std::array buffer; -+ ssize_t bytes_read; -+ -+ while ((bytes_read = read(stdout_pipe[0], buffer.data(), buffer.size())) > 0) { -+ stdout_str.append(buffer.data(), bytes_read); -+ } -+ -+ while ((bytes_read = read(stderr_pipe[0], buffer.data(), buffer.size())) > 0) { -+ stderr_str.append(buffer.data(), bytes_read); -+ } -+ -+ close(stdout_pipe[0]); -+ close(stderr_pipe[0]); -+ waitpid(pid, nullptr, 0); -+ } -+#endif -+} -+ -+bool directory_exists(const std::string& path) { -+ struct stat info; -+ if (stat(path.c_str(), &info) != 0) { -+ return false; // Path doesn't exist or can't be accessed -+ } -+ return (info.st_mode & S_IFDIR) != 0; // Check if it is a directory -+} -+ -+bool create_directory(const std::string& path) { -+#ifdef _WIN32 -+ return _mkdir(path.c_str()) == 0 || errno == EEXIST; // EEXIST means the directory already exists -+#else -+ return mkdir(path.c_str(), 0755) == 0 || errno == EEXIST; // 0755 is the directory permissions -+#endif -+} -+ -+std::string to_uppercase(const std::string& input) { -+ std::string result = input; -+ for (char& c : result) { -+ c = std::toupper(c); -+ } -+ return result; -+} -+ -+bool string_ends_with(const std::string& str, const std::string& suffix) { -+ if (suffix.size() > str.size()) { -+ return false; -+ } -+ return std::equal(suffix.rbegin(), suffix.rend(), str.rbegin()); -+} -+ -+static const char path_separator = '/'; -+ -+std::string join_paths(const std::string& path1, const std::string& path2) { -+ return path1 + path_separator + path2; -+} -+ -+std::string basename(const std::string &path) { -+ return path.substr(path.find_last_of("/\\") + 1); -+} -+ -+// variables to track number of compiles in progress -+static uint32_t compile_count = 0; -+static std::mutex compile_count_mutex; -+static std::condition_variable compile_count_cond; -+ -+void string_to_spv_func(const std::string& _name, const std::string& in_fname, const std::map& defines, bool fp16 = true, bool coopmat = false, bool coopmat2 = false, bool f16acc = false) { -+ std::string name = _name + (f16acc ? "_f16acc" : "") + (coopmat ? "_coopmat" : "") + (coopmat2 ? "_cm2" : (fp16 ? "" : "_fp32")); -+ std::string out_fname = join_paths(output_dir, name + ".spv"); -+ std::string in_path = join_paths(input_dir, in_fname); -+ -+ std::string target_env = (name.find("_cm2") != std::string::npos) ? "--target-env=vulkan1.3" : "--target-env=vulkan1.2"; -+ -+ // disable spirv-opt for coopmat shaders for https://github.com/ggerganov/llama.cpp/issues/10734 -+ std::string opt_level = coopmat ? "" : "-O"; -+ -+ #ifdef _WIN32 -+ std::vector cmd = {GLSLC, "-fshader-stage=compute", target_env, opt_level, "\"" + in_path + "\"", "-o", "\"" + out_fname + "\""}; -+ #else -+ std::vector cmd = {GLSLC, "-fshader-stage=compute", target_env, opt_level, in_path, "-o", out_fname}; -+ #endif -+ -+ #ifdef GGML_VULKAN_SHADER_DEBUG_INFO -+ cmd.push_back("-g"); -+ #endif -+ -+ for (const auto& define : defines) { -+ cmd.push_back("-D" + define.first + "=" + define.second); -+ } -+ -+ std::string command; -+ for (const auto& part : cmd) { -+ command += part + " "; -+ } -+ -+ std::string stdout_str, stderr_str; -+ try { -+ // std::cout << "Executing command: "; -+ // for (const auto& part : cmd) { -+ // std::cout << part << " "; -+ // } -+ // std::cout << std::endl; -+ -+ execute_command(command, stdout_str, stderr_str); -+ if (!stderr_str.empty()) { -+ std::cerr << "cannot compile " << name << "\n\n" << command << "\n\n" << stderr_str << std::endl; -+ return; -+ } -+ -+ std::lock_guard guard(lock); -+ shader_fnames.push_back(std::make_pair(name, out_fname)); -+ } catch (const std::exception& e) { -+ std::cerr << "Error executing command for " << name << ": " << e.what() << std::endl; -+ } -+ { -+ std::lock_guard guard(compile_count_mutex); -+ assert(compile_count > 0); -+ compile_count--; -+ } -+ compile_count_cond.notify_all(); -+} -+ -+std::map merge_maps(const std::map& a, const std::map& b) { -+ std::map result = a; -+ result.insert(b.begin(), b.end()); -+ return result; -+} -+ -+static std::vector> compiles; -+void string_to_spv(const std::string& _name, const std::string& in_fname, const std::map& defines, bool fp16 = true, bool coopmat = false, bool coopmat2 = false, bool f16acc = false) { -+ { -+ // wait until fewer than N compiles are in progress. -+ // 16 is an arbitrary limit, the goal is to avoid "failed to create pipe" errors. -+ uint32_t N = 16; -+ std::unique_lock guard(compile_count_mutex); -+ while (compile_count >= N) { -+ compile_count_cond.wait(guard); -+ } -+ compile_count++; -+ } -+ compiles.push_back(std::async(string_to_spv_func, _name, in_fname, defines, fp16, coopmat, coopmat2, f16acc)); -+} -+ -+void matmul_shaders(bool fp16, bool matmul_id, bool coopmat, bool coopmat2, bool f16acc) { -+ std::string load_vec = coopmat2 ? "1" : fp16 ? "8" : "4"; -+ std::string aligned_b_type_f32 = coopmat2 ? "float" : fp16 ? "mat2x4" : "vec4"; -+ std::string aligned_b_type_f16 = coopmat2 ? "float16_t" : fp16 ? "f16mat2x4" : "f16vec4"; -+ -+ std::map base_dict = {{"FLOAT_TYPE", (coopmat2 || fp16) ? "float16_t" : "float"}}; -+ std::string shader_name = "matmul"; -+ -+ if (matmul_id) { -+ base_dict["MUL_MAT_ID"] = "1"; -+ shader_name = "matmul_id"; -+ } -+ -+ if (fp16) { -+ base_dict["FLOAT16"] = "1"; -+ } -+ -+ base_dict["ACC_TYPE"] = f16acc ? "float16_t" : "float"; -+ -+ if (coopmat) { -+ base_dict["COOPMAT"] = "1"; -+ } -+ -+ base_dict["ACC_TYPE"] = f16acc ? "float16_t" : "float"; -+ -+ std::string source_name = coopmat2 ? "mul_mm_cm2.comp" : "mul_mm.comp"; -+ -+ // Shaders with f16 B_TYPE -+ string_to_spv(shader_name + "_f32_f16", source_name, merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}, }), fp16, coopmat, coopmat2, f16acc); -+ string_to_spv(shader_name + "_f32_f16_aligned", source_name, merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); -+ -+ string_to_spv(shader_name + "_f16_aligned", source_name, merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); -+ string_to_spv(shader_name + "_f16", source_name, merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); -+ -+ for (const auto& tname : type_names) { -+ std::string data_a_key = "DATA_A_" + to_uppercase(tname); -+ // For unaligned, load one at a time for f32/f16, or two at a time for quants -+ std::string load_vec_a_unaligned = (coopmat2 || tname == "f32" || tname == "f16") ? "1" : "2"; -+ // For aligned matmul loads -+ std::string load_vec_a = (coopmat2 || tname == "f32" || tname == "f16") ? load_vec : "2"; -+ -+ string_to_spv(shader_name + "_" + tname + "_f32", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}}), fp16, coopmat, coopmat2, f16acc); -+ string_to_spv(shader_name + "_" + tname + "_f32_aligned", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); -+ -+ if (tname != "f16" && tname != "f32") { -+ string_to_spv(shader_name + "_" + tname + "_f16", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}}), fp16, coopmat, coopmat2, f16acc); -+ string_to_spv(shader_name + "_" + tname + "_f16_aligned", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); -+ } -+ } -+} -+ -+void process_shaders() { -+ std::cout << "ggml_vulkan: Generating and compiling shaders to SPIR-V" << std::endl; -+ std::map base_dict = {{"FLOAT_TYPE", "float"}}; -+ -+ // matmul -+ for (const auto& matmul_id : {false, true}) { -+ // No coopmats -+ // fp32 -+ matmul_shaders(false, matmul_id, false, false, false); -+ -+ // fp16, fp32acc and fp16acc -+ matmul_shaders(true, matmul_id, false, false, false); -+ matmul_shaders(true, matmul_id, false, false, true); -+ -+ // Coopmat, fp32acc and fp16acc -+ matmul_shaders(true, matmul_id, true, false, false); -+ matmul_shaders(true, matmul_id, true, false, true); -+ -+#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) -+ // Coopmat2, fp32acc and fp16acc -+ matmul_shaders(true, matmul_id, false, true, false); -+ matmul_shaders(true, matmul_id, false, true, true); -+#endif -+ } -+ -+#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) -+ // flash attention -+ for (const auto& f16acc : {false, true}) { -+ std::string acctype = f16acc ? "float16_t" : "float"; -+ -+ for (const auto& tname : type_names) { -+ if (tname == "f32") { -+ continue; -+ } -+ -+ if (tname == "f16") { -+ string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm2.comp", -+ merge_maps(base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"ACC_TYPE", acctype}}), true, false, true, f16acc); -+ } else { -+ std::string data_a_key = "DATA_A_" + to_uppercase(tname); -+ string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm2.comp", -+ merge_maps(base_dict, {{data_a_key, "1"}, {"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"ACC_TYPE", acctype}, {"DEQUANTFUNC", "dequantFunc"+to_uppercase(tname) }, {"BLOCK_SIZE", "QUANT_K_"+to_uppercase(tname) }}), true, false, true, f16acc); -+ } -+ } -+ } -+#endif -+ -+ for (const auto& tname : type_names) { -+ // mul mat vec -+ std::string data_a_key = "DATA_A_" + to_uppercase(tname); -+ std::string shader = (string_ends_with(tname, "_k")) ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp"; -+ -+ string_to_spv("mul_mat_vec_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPE_VEC2", "vec2"}, {"B_TYPE_VEC4", "vec4"}, {"D_TYPE", "float"}})); -+ string_to_spv("mul_mat_vec_" + tname + "_f16_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"B_TYPE_VEC2", "f16vec2"}, {"B_TYPE_VEC4", "f16vec4"}, {"D_TYPE", "float"}})); -+ -+ string_to_spv("mul_mat_vec_id_" + tname + "_f32", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPE_VEC2", "vec2"}, {"B_TYPE_VEC4", "vec4"}, {"D_TYPE", "float"}})); -+ -+ // Dequant shaders -+ if (tname != "f16") { -+ string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}})); -+ } -+ -+ if (!string_ends_with(tname, "_k")) { -+ shader = (tname == "f32" || tname == "f16") ? "get_rows.comp" : "get_rows_quant.comp"; -+ -+ if (tname == "f16") { -+ string_to_spv("get_rows_" + tname, shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}})); -+ } else { -+ string_to_spv("get_rows_" + tname, shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}})); -+ } -+ string_to_spv("get_rows_" + tname + "_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float"}})); -+ } -+ } -+ -+ string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); -+ string_to_spv("mul_mat_vec_nc_f16_f32", "mul_mat_vec_nc.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); -+ -+ // Norms -+ string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); -+ string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); -+ string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); -+ -+ string_to_spv("cpy_f32_f32", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ string_to_spv("cpy_f32_f16", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); -+ string_to_spv("cpy_f16_f16", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); -+ string_to_spv("contig_cpy_f32_f32", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ string_to_spv("contig_cpy_f32_f16", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); -+ string_to_spv("contig_cpy_f16_f16", "contig_copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); -+ -+ string_to_spv("add_f32", "add.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); -+ string_to_spv("add_f16_f32_f16", "add.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"FLOAT_TYPE", "float"}}); -+ -+ string_to_spv("acc_f32", "acc.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); -+ -+ string_to_spv("split_k_reduce", "mul_mat_split_k_reduce.comp", {}); -+ -+ string_to_spv("mul_f32", "mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); -+ -+ string_to_spv("div_f32", "div.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); -+ -+ string_to_spv("repeat_f32", "repeat.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ -+ string_to_spv("scale_f32", "scale.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); -+ -+ string_to_spv("sqr_f32", "square.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); -+ -+ string_to_spv("sin_f32", "sin.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); -+ -+ string_to_spv("cos_f32", "cos.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); -+ -+ string_to_spv("clamp_f32", "clamp.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); -+ -+ string_to_spv("pad_f32", "pad.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ -+ string_to_spv("concat_f32", "concat.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); -+ string_to_spv("concat_f16", "concat.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); -+ string_to_spv("concat_i32", "concat.comp", {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}}); -+ -+ string_to_spv("upscale_f32", "upscale.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); -+ -+ string_to_spv("gelu_f32", "gelu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ string_to_spv("gelu_quick_f32", "gelu_quick.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ string_to_spv("silu_f32", "silu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ string_to_spv("relu_f32", "relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ string_to_spv("leaky_relu_f32", "leaky_relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ string_to_spv("tanh_f32", "tanh.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ -+ string_to_spv("diag_mask_inf_f32", "diag_mask_inf.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ -+ string_to_spv("soft_max_f32", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); -+ string_to_spv("soft_max_f32_f16", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}})); -+ -+ string_to_spv("rope_norm_f32", "rope_norm.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ string_to_spv("rope_norm_f16", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); -+ string_to_spv("rope_norm_f16_rte", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}}); -+ -+ string_to_spv("rope_neox_f32", "rope_neox.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); -+ string_to_spv("rope_neox_f16", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); -+ string_to_spv("rope_neox_f16_rte", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}}); -+ -+ string_to_spv("argsort_f32", "argsort.comp", {{"A_TYPE", "float"}}); -+ -+ string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); -+ -+ string_to_spv("im2col_f32", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); -+ string_to_spv("im2col_f32_f16", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}})); -+ string_to_spv("im2col_f32_f16_rte", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}})); -+ -+ string_to_spv("timestep_embedding_f32", "timestep_embedding.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); -+ -+ string_to_spv("pool2d_f32", "pool2d.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); -+ -+ string_to_spv("rwkv_wkv6_f32", "wkv6.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); -+ -+ for (auto &c : compiles) { -+ c.wait(); -+ } -+} -+ -+void write_output_files() { -+ FILE* hdr = fopen(target_hpp.c_str(), "w"); -+ FILE* src = fopen(target_cpp.c_str(), "w"); -+ -+ fprintf(hdr, "#include \n\n"); -+ fprintf(src, "#include \"%s\"\n\n", basename(target_hpp).c_str()); -+ -+ for (const auto& pair : shader_fnames) { -+ const std::string& name = pair.first; -+ #ifdef _WIN32 -+ std::string path = pair.second; -+ std::replace(path.begin(), path.end(), '/', '\\' ); -+ #else -+ const std::string& path = pair.second; -+ #endif -+ -+ FILE* spv = fopen(path.c_str(), "rb"); -+ if (!spv) { -+ std::cerr << "Error opening SPIR-V file: " << path << " (" << strerror(errno) << ")\n"; -+ continue; -+ } -+ -+ fseek(spv, 0, SEEK_END); -+ size_t size = ftell(spv); -+ fseek(spv, 0, SEEK_SET); -+ -+ std::vector data(size); -+ size_t read_size = fread(data.data(), 1, size, spv); -+ fclose(spv); -+ if (read_size != size) { -+ std::cerr << "Error reading SPIR-V file: " << path << " (" << strerror(errno) << ")\n"; -+ continue; -+ } -+ -+ fprintf(hdr, "extern unsigned char %s_data[%zu];\n", name.c_str(), size); -+ fprintf(hdr, "const uint64_t %s_len = %zu;\n\n", name.c_str(), size); -+ -+ fprintf(src, "unsigned char %s_data[%zu] = {\n", name.c_str(), size); -+ for (size_t i = 0; i < size; ++i) { -+ fprintf(src, "0x%02x,", data[i]); -+ if ((i + 1) % 12 == 0) fprintf(src, "\n"); -+ } -+ fprintf(src, "\n};\n\n"); -+ -+ if (!no_clean) { -+ std::remove(path.c_str()); -+ } -+ } -+ -+ fclose(hdr); -+ fclose(src); -+} -+} -+ -+int main(int argc, char** argv) { -+ std::map args; -+ for (int i = 1; i < argc; ++i) { -+ std::string arg = argv[i]; -+ if (arg.rfind("--", 0) == 0) { -+ if (i + 1 < argc && argv[i + 1][0] != '-') { -+ args[arg] = argv[i + 1]; -+ ++i; -+ } else { -+ args[arg] = ""; -+ } -+ } -+ } -+ -+ if (args.find("--glslc") != args.end()) { -+ GLSLC = args["--glslc"]; // Path to glslc -+ } -+ if (args.find("--input-dir") != args.end()) { -+ input_dir = args["--input-dir"]; // Directory containing shader sources -+ } -+ if (args.find("--output-dir") != args.end()) { -+ output_dir = args["--output-dir"]; // Directory for containing SPIR-V output -+ } -+ if (args.find("--target-hpp") != args.end()) { -+ target_hpp = args["--target-hpp"]; // Path to generated header file -+ } -+ if (args.find("--target-cpp") != args.end()) { -+ target_cpp = args["--target-cpp"]; // Path to generated cpp file -+ } -+ if (args.find("--no-clean") != args.end()) { -+ no_clean = true; // Keep temporary SPIR-V files in output-dir after build -+ } -+ -+ if (!directory_exists(input_dir)) { -+ std::cerr << "\"" << input_dir << "\" must be a valid directory containing shader sources" << std::endl; -+ return EXIT_FAILURE; -+ } -+ -+ if (!directory_exists(output_dir)) { -+ if (!create_directory(output_dir)) { -+ std::cerr << "Error creating output directory: " << output_dir << "\n"; -+ return EXIT_FAILURE; -+ } -+ } -+ -+ process_shaders(); -+ -+ write_output_files(); -+ -+ return EXIT_SUCCESS; -+} -diff --git a/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp -new file mode 100644 -index 00000000..35cc6c45 ---- /dev/null -+++ b/ml/backend/ggml/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp -@@ -0,0 +1,87 @@ -+#version 450 -+ -+#extension GL_EXT_control_flow_attributes : require -+ -+#define BLOCK_SIZE 64 -+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; -+ -+layout(push_constant) uniform Parameters { -+ uint B; -+ uint T; -+ uint C; -+ uint H; -+}; -+ -+layout(binding = 0) readonly buffer KBuf { A_TYPE k[]; }; -+layout(binding = 1) readonly buffer VBuf { A_TYPE v[]; }; -+layout(binding = 2) readonly buffer RBuf { A_TYPE r[]; }; -+layout(binding = 3) readonly buffer TimeFBuf { A_TYPE tf[]; }; -+layout(binding = 4) readonly buffer TimeDBuf { A_TYPE td[]; }; -+layout(binding = 5) readonly buffer StateBuf { A_TYPE state_in[]; }; -+layout(binding = 6) buffer DstBuf { A_TYPE dst[]; }; -+ -+shared A_TYPE _k[BLOCK_SIZE], _r[BLOCK_SIZE], _tf[BLOCK_SIZE], _td[BLOCK_SIZE]; -+ -+void main() { -+ const uint head_size = BLOCK_SIZE; -+ const uint batch_id = gl_WorkGroupID.x / H; -+ const uint head_id = gl_WorkGroupID.x % H; -+ const uint tid = gl_LocalInvocationID.x; -+ -+ const uint state_size = C * head_size; -+ const uint n_seq_tokens = T / B; -+ -+ if (batch_id >= B || head_id >= H) { -+ return; -+ } -+ -+ A_TYPE state[BLOCK_SIZE]; -+ [[unroll]] for (uint i = 0; i < head_size; i++) { -+ state[i] = state_in[batch_id * state_size + head_id * head_size * head_size -+ + i * head_size + tid]; -+ } -+ -+ barrier(); -+ _tf[tid] = tf[head_id * head_size + tid]; -+ barrier(); -+ -+ const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid; -+ const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid; -+ -+ for (uint t = start_t; t < end_t; t += C) { -+ barrier(); -+ _k[tid] = k[t]; -+ _r[tid] = r[t]; -+ _td[tid] = td[t]; -+ barrier(); -+ -+ const A_TYPE v_val = v[t]; -+ A_TYPE y = 0.0; -+ -+ [[unroll]] for (uint j = 0; j < head_size; j += 4) { -+ vec4 k_vec = vec4(_k[j], _k[j+1], _k[j+2], _k[j+3]); -+ vec4 r_vec = vec4(_r[j], _r[j+1], _r[j+2], _r[j+3]); -+ vec4 tf_vec = vec4(_tf[j], _tf[j+1], _tf[j+2], _tf[j+3]); -+ vec4 td_vec = vec4(_td[j], _td[j+1], _td[j+2], _td[j+3]); -+ vec4 s_vec = vec4(state[j], state[j+1], state[j+2], state[j+3]); -+ -+ vec4 kv = k_vec * v_val; -+ -+ vec4 temp = tf_vec * kv + s_vec; -+ y += dot(r_vec, temp); -+ -+ s_vec = s_vec * td_vec + kv; -+ state[j] = s_vec.x; -+ state[j+1] = s_vec.y; -+ state[j+2] = s_vec.z; -+ state[j+3] = s_vec.w; -+ } -+ -+ dst[t] = y; -+ } -+ -+ [[unroll]] for (uint i = 0; i < head_size; i++) { -+ dst[T * C + batch_id * state_size + head_id * head_size * head_size -+ + i * head_size + tid] = state[i]; -+ } -+} --- -2.43.0 -