KBaba7/llama.cpp
0
1#include "ggml-impl.h"2#include "ggml-backend.h"3#include "ggml-backend-impl.h"4#include "ggml-kompute.h"5 6// These are generated at build time by cmake custom command7#include "shaderop_scale.h"8#include "shaderop_scale_8.h"9#include "shaderop_add.h"10#include "shaderop_addrow.h"11#include "shaderop_mul.h"12#include "shaderop_silu.h"13#include "shaderop_relu.h"14#include "shaderop_gelu.h"15#include "shaderop_softmax.h"16#include "shaderop_norm.h"17#include "shaderop_rmsnorm.h"18#include "shaderop_diagmask.h"19#include "shaderop_mul_mat_f16.h"20#include "shaderop_mul_mat_q8_0.h"21#include "shaderop_mul_mat_q4_0.h"22#include "shaderop_mul_mat_q4_1.h"23#include "shaderop_mul_mat_q4_k.h"24#include "shaderop_mul_mat_q6_k.h"25#include "shaderop_mul_mat_mat_f32.h"26#include "shaderop_getrows_f32.h"27#include "shaderop_getrows_f16.h"28#include "shaderop_getrows_q4_0.h"29#include "shaderop_getrows_q4_1.h"30#include "shaderop_getrows_q6_k.h"31#include "shaderop_rope_norm_f16.h"32#include "shaderop_rope_norm_f32.h"33#include "shaderop_rope_neox_f16.h"34#include "shaderop_rope_neox_f32.h"35#include "shaderop_cpy_f16_f16.h"36#include "shaderop_cpy_f16_f32.h"37#include "shaderop_cpy_f32_f16.h"38#include "shaderop_cpy_f32_f32.h"39 40#include <algorithm>41#include <array>42#include <cassert>43#include <cstdint>44#include <cstdio>45#include <cstring>46#include <iostream>47#include <memory>48#include <mutex>49#include <stdexcept>50#include <string>51#include <unordered_map>52#include <utility>53#include <vector>54 55#include <kompute/Kompute.hpp>56#include <vulkan/vulkan.hpp>57 58#ifdef __linux__59#include <cstdlib> // for setenv60#endif61 62#define QK4_0 3263#define QR4_0 264#define QK4_1 3265#define QK_NL 1666 67typedef ggml_fp16_t half;68 69static std::string ggml_kompute_format_name(int device) {70 return "Kompute" + std::to_string(device);71}72 73struct ggml_kompute_context {74 int device;75 std::string name;76 std::shared_ptr<vk::DescriptorPool> pool;77 78 ggml_kompute_context(int device)79 : device(device), name(ggml_kompute_format_name(device)) {}80};81 82// FIXME: It would be good to consolidate the kompute manager and the kompute context into one object83// and consolidate the init functions and simplify object lifetime management. As it currently stands,84// we *have* to have the kompute manager no matter what for device discovery, but the kompute context85// is only created when a device is set and vulkan is explicitly turned on.86static ggml_kompute_context *s_kompute_context = nullptr;87 88class kompute_manager {89 kp::Manager *s_mgr = nullptr;90 91public:92 kp::Manager *operator()() {93 if (s_mgr && !s_mgr->hasInstance()) {94 destroy();95 }96 if (!s_mgr) {97 s_mgr = new kp::Manager;98 }99 return s_mgr;100 }101 102 void destroy() {103 delete s_mgr;104 s_mgr = nullptr;105 }106};107 108static kompute_manager komputeManager;109 110struct ggml_vk_memory {111 void *data = nullptr;112 size_t size = 0;113 vk::DeviceMemory *primaryMemory = nullptr;114 vk::Buffer *primaryBuffer = nullptr;115 vk::DeviceMemory *stagingMemory = nullptr;116 vk::Buffer *stagingBuffer = nullptr;117};118 119#ifdef __linux__120__attribute__((constructor))121static void enable_sam() {122 setenv("RADV_PERFTEST", "sam", false);123}124#endif125 126static bool ggml_vk_checkPhysicalDeviceFeatures(vk::PhysicalDevice physical_device) {127 vk::PhysicalDeviceFeatures availableFeatures;128 physical_device.getFeatures(&availableFeatures);129 130 if (!availableFeatures.shaderInt16)131 return false;132 133 vk::PhysicalDeviceVulkan11Features availableFeatures11;134 vk::PhysicalDeviceVulkan12Features availableFeatures12;135 136 availableFeatures11.pNext = &availableFeatures12;137 availableFeatures12.pNext = nullptr;138 139 vk::PhysicalDeviceFeatures2 features2;140 features2.pNext = &availableFeatures11;141 142 physical_device.getFeatures2(&features2);143 144 if (!availableFeatures11.uniformAndStorageBuffer16BitAccess ||145 !availableFeatures11.storageBuffer16BitAccess) {146 return false;147 }148 149 if (!availableFeatures12.storageBuffer8BitAccess ||150 !availableFeatures12.uniformAndStorageBuffer8BitAccess ||151 !availableFeatures12.shaderFloat16 ||152 !availableFeatures12.shaderInt8) {153 return false;154 }155 156 return true;157}158 159static const char * ggml_vk_getVendorName(uint32_t vendorID) {160 switch (vendorID) {161 case 0x10DE:162 return "nvidia";163 case 0x1002:164 return "amd";165 case 0x8086:166 return "intel";167 default:168 return "unknown";169 }170}171 172static std::vector<ggml_vk_device> ggml_vk_available_devices_internal(size_t memoryRequired) {173 std::vector<ggml_vk_device> results;174 if (!komputeManager()->hasVulkan() || !komputeManager()->hasInstance())175 return results;176 177 std::vector<vk::PhysicalDevice> physical_devices;178 try {179 physical_devices = komputeManager()->listDevices();180 } catch (vk::SystemError & err) {181 std::cerr << __func__ << ": ignoring Vulkan exception: " << err.what() << "\n";182 return results;183 }184 185 uint32_t deviceCount = physical_devices.size();186 if (deviceCount == 0)187 return results;188 189 std::unordered_map<std::string, size_t> count_by_name;190 191 for (uint32_t i = 0; i < deviceCount; i++) {192 const auto & physical_device = physical_devices[i];193 194 VkPhysicalDeviceProperties dev_props = physical_device.getProperties();195 VkPhysicalDeviceMemoryProperties memoryProperties = physical_device.getMemoryProperties();196 const uint32_t major = VK_VERSION_MAJOR(dev_props.apiVersion);197 const uint32_t minor = VK_VERSION_MINOR(dev_props.apiVersion);198 if (major < 1 || minor < 2)199 continue;200 201 if (!ggml_vk_checkPhysicalDeviceFeatures(physical_device))202 continue;203 204 size_t heapSize = 0;205 for (uint32_t j = 0; j < memoryProperties.memoryHeapCount; ++j) {206 VkMemoryHeap heap = memoryProperties.memoryHeaps[j];207 if (heap.flags & VK_MEMORY_HEAP_DEVICE_LOCAL_BIT) {208 heapSize = heap.size;209 break;210 }211 }212 213 if (heapSize < memoryRequired)214 continue;215 216 auto ext_props = physical_device.enumerateDeviceExtensionProperties();217 bool has_maintenance4 = false;218 219 // Check if maintenance4 is supported220 for (const auto & properties : ext_props) {221 if (strcmp("VK_KHR_maintenance4", properties.extensionName) == 0) {222 has_maintenance4 = true;223 }224 }225 226 vk::PhysicalDeviceSubgroupProperties subgroup_props;227 vk::PhysicalDeviceProperties2 dev_props2;228 vk::PhysicalDeviceMaintenance3Properties dev_props3;229 vk::PhysicalDeviceMaintenance4Properties dev_props4;230 dev_props2.pNext = &dev_props3;231 dev_props3.pNext = &subgroup_props;232 if (has_maintenance4) {233 subgroup_props.pNext = &dev_props4;234 }235 physical_device.getProperties2(&dev_props2);236 237 if (subgroup_props.subgroupSize < 32)238 continue;239 240 ggml_vk_device d;241 d.index = i;242 d.type = dev_props.deviceType;243 d.heapSize = heapSize;244 d.vendor = strdup(ggml_vk_getVendorName(dev_props.vendorID));245 d.subgroupSize = subgroup_props.subgroupSize;246 d.bufferAlignment = dev_props.limits.minStorageBufferOffsetAlignment;247 248 if (has_maintenance4) {249 d.maxAlloc = std::min(dev_props3.maxMemoryAllocationSize, dev_props4.maxBufferSize);250 } else {251 d.maxAlloc = dev_props3.maxMemoryAllocationSize;252 }253 254 std::string name(dev_props.deviceName);255 size_t n_idx = ++count_by_name[name];256 if (n_idx > 1) {257 name += " (" + std::to_string(n_idx) + ")";258 }259 d.name = strdup(name.c_str());260 261 results.push_back(d);262 }263 264 std::stable_sort(results.begin(), results.end(),265 [](const ggml_vk_device& lhs, const ggml_vk_device& rhs) -> bool {266 if (lhs.type != rhs.type) {267 if (lhs.type == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) return true;268 if (rhs.type == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) return false;269 270 if (lhs.type == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU) return true;271 if (rhs.type == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU) return false;272 }273 return lhs.heapSize < rhs.heapSize;274 }275 );276 277 return results;278}279 280static std::vector<ggml_vk_device>& ggml_vk_available_devices() {281 static std::vector<ggml_vk_device> devices = ggml_vk_available_devices_internal(0);282 return devices;283}284 285static void ggml_vk_filterByVendor(std::vector<ggml_vk_device>& devices, const std::string& targetVendor) {286 devices.erase(287 std::remove_if(devices.begin(), devices.end(),288 [&targetVendor](const ggml_vk_device& device) {289 return device.vendor != targetVendor;290 }),291 devices.end()292 );293}294 295static void ggml_vk_filterByName(std::vector<ggml_vk_device>& devices, const std::string& targetName) {296 devices.erase(297 std::remove_if(devices.begin(), devices.end(),298 [&targetName](const ggml_vk_device& device) {299 return device.name != targetName;300 }),301 devices.end()302 );303}304 305static bool ggml_vk_get_device(ggml_vk_device * device, size_t memoryRequired, const std::string & name) {306 if (name.empty())307 return false;308 309 auto devices = ggml_vk_available_devices_internal(memoryRequired);310 if (name == "amd" || name == "nvidia" || name == "intel") {311 ggml_vk_filterByVendor(devices, name);312 } else if (name != "gpu") {313 ggml_vk_filterByName(devices, name);314 }315 316 if (devices.empty())317 return false;318 319 *device = devices.front();320 return true;321}322 323bool ggml_vk_get_device(ggml_vk_device * device, size_t memoryRequired, const char * name) {324 return ggml_vk_get_device(device, memoryRequired, std::string(name));325}326 327bool ggml_vk_has_vulkan() {328 return komputeManager()->hasVulkan();329}330 331bool ggml_vk_has_device() {332 return komputeManager()->hasDevice();333}334 335ggml_vk_device ggml_vk_current_device() {336 if (!komputeManager()->hasDevice())337 return ggml_vk_device();338 339 auto devices = ggml_vk_available_devices();340 ggml_vk_filterByName(devices, komputeManager()->physicalDevice()->getProperties().deviceName.data());341 GGML_ASSERT(!devices.empty());342 return devices.front();343}344 345static346void ggml_vk_allocate_descriptor_pool(struct ggml_kompute_context * ctx, size_t size) {347 std::vector<vk::DescriptorPoolSize> descriptorPoolSizes = {348 vk::DescriptorPoolSize(349 vk::DescriptorType::eStorageBuffer,350 4 * size // Descriptor count is number of possible tensors to pass into an algorithm351 )352 };353 354 vk::DescriptorPoolCreateInfo descriptorPoolInfo(355 vk::DescriptorPoolCreateFlags(),356 size, // Max sets357 static_cast<uint32_t>(descriptorPoolSizes.size()),358 descriptorPoolSizes.data());359 360 ctx->pool = std::make_shared<vk::DescriptorPool>();361 vk::Result r = komputeManager()->device()->createDescriptorPool(362 &descriptorPoolInfo, nullptr, ctx->pool.get());363 if (r != vk::Result::eSuccess)364 std::cerr << "Error allocating descriptor pool" << vk::to_string(r);365}366 367static368void ggml_vk_free_descriptor_pool(struct ggml_kompute_context * ctx) {369 if (ctx->pool) {370 komputeManager()->device()->destroy(371 *ctx->pool,372 (vk::Optional<const vk::AllocationCallbacks>)nullptr);373 ctx->pool = nullptr;374 }375}376 377static378vk::Buffer *ggml_vk_allocate_buffer(size_t size) {379 vk::BufferCreateInfo bufferCreateInfo;380 bufferCreateInfo.size = size;381 bufferCreateInfo.usage = vk::BufferUsageFlagBits::eStorageBuffer |382 vk::BufferUsageFlagBits::eTransferSrc |383 vk::BufferUsageFlagBits::eTransferDst;384 bufferCreateInfo.sharingMode = vk::SharingMode::eExclusive;385 386 vk::Buffer *vkBuffer = new vk::Buffer;387 vk::Result r = komputeManager()->device()->createBuffer(&bufferCreateInfo, nullptr, vkBuffer);388 if (r != vk::Result::eSuccess)389 std::cerr << "Error allocating buffer " << vk::to_string(r) << std::endl;390 return vkBuffer;391}392 393static394vk::DeviceMemory *ggml_vk_allocate(size_t size, vk::MemoryPropertyFlags flags, vk::MemoryRequirements requirements, bool *isHostVisible) {395 396 uint32_t memoryTypeIndex = -1;397 bool memoryTypeIndexFound = false;398 vk::PhysicalDeviceMemoryProperties memoryProperties = komputeManager()->physicalDevice()->getMemoryProperties();399 for (uint32_t i = 0; i < memoryProperties.memoryTypeCount; i++) {400 const vk::MemoryType &memoryType = memoryProperties.memoryTypes[i];401 const vk::MemoryHeap &memoryHeap = memoryProperties.memoryHeaps[memoryType.heapIndex];402 if (memoryHeap.size < size) {403 continue;404 }405 406 if (requirements.memoryTypeBits & (1 << i)) {407 if (((memoryProperties.memoryTypes[i]).propertyFlags &408 flags) == flags) {409 memoryTypeIndex = i;410 memoryTypeIndexFound = true;411 if (isHostVisible && (memoryProperties.memoryTypes[i].propertyFlags & vk::MemoryPropertyFlagBits::eHostVisible)) {412 *isHostVisible = true;413 }414 break;415 }416 }417 }418 if (!memoryTypeIndexFound) {419 throw std::runtime_error(420 "Memory type index for buffer creation not found");421 }422 423 vk::MemoryAllocateInfo allocInfo;424 allocInfo.allocationSize = size;425 allocInfo.memoryTypeIndex = memoryTypeIndex;426 vk::DeviceMemory *vkDeviceMemory = new vk::DeviceMemory;427 vk::Result r = komputeManager()->device()->allocateMemory(&allocInfo, nullptr, vkDeviceMemory);428 if (r != vk::Result::eSuccess) {429 std::cerr << "Error allocating memory " << vk::to_string(r) << std::endl;430 throw std::runtime_error("Error allocating vulkan memory.");431 }432 return vkDeviceMemory;433}434 435static size_t ggml_vk_aligned_offset(ggml_backend_buffer_t buffer, size_t offset) {436 size_t minStorageBufferOffsetAlignment = ggml_backend_buffer_get_alignment(buffer);437 438 // If offset is already aligned, return it directly439 if (offset % minStorageBufferOffsetAlignment == 0) {440 return offset;441 }442 443 // Otherwise, return the largest multiple of minStorageBufferOffsetAlignment less than offset444 return (offset / minStorageBufferOffsetAlignment) * minStorageBufferOffsetAlignment;445}446 447static ggml_vk_memory ggml_vk_allocate(size_t size) {448 ggml_vk_memory memory;449 bool isHostVisible = false;450 {451 memory.primaryBuffer = ggml_vk_allocate_buffer(size);452 vk::MemoryRequirements memoryRequirements = komputeManager()->device()->getBufferMemoryRequirements(*memory.primaryBuffer);453 vk::MemoryPropertyFlags memoryPropertyFlags = vk::MemoryPropertyFlagBits::eDeviceLocal;454 memory.primaryMemory = ggml_vk_allocate(size, memoryPropertyFlags, memoryRequirements, &isHostVisible);455 komputeManager()->device()->bindBufferMemory(*memory.primaryBuffer, *memory.primaryMemory, 0);456 if (isHostVisible) {457 vk::Result r = komputeManager()->device()->mapMemory(*memory.primaryMemory, 0, size, vk::MemoryMapFlags(), &memory.data);458 if (r != vk::Result::eSuccess)459 std::cerr << "Error mapping memory" << vk::to_string(r);460 }461 }462 463 if (!isHostVisible) {464 memory.stagingBuffer = ggml_vk_allocate_buffer(size);465 vk::MemoryRequirements memoryRequirements = komputeManager()->device()->getBufferMemoryRequirements(*memory.stagingBuffer);466 vk::MemoryPropertyFlags memoryPropertyFlags = vk::MemoryPropertyFlagBits::eHostVisible |467 vk::MemoryPropertyFlagBits::eHostCoherent |468 vk::MemoryPropertyFlagBits::eHostCached;469 memory.stagingMemory = ggml_vk_allocate(size, memoryPropertyFlags, memoryRequirements, &isHostVisible);470 komputeManager()->device()->bindBufferMemory(*memory.stagingBuffer, *memory.stagingMemory, 0);471 vk::Result r = komputeManager()->device()->mapMemory(*memory.stagingMemory, 0, size, vk::MemoryMapFlags(), &memory.data);472 if (r != vk::Result::eSuccess)473 std::cerr << "Error mapping memory" << vk::to_string(r);474 }475 476 memory.size = size;477 return memory;478}479 480static void ggml_vk_free_memory(ggml_vk_memory &memory)481{482 komputeManager()->device()->destroy(483 *memory.primaryBuffer,484 (vk::Optional<const vk::AllocationCallbacks>)nullptr);485 if (memory.stagingBuffer) {486 komputeManager()->device()->destroy(487 *memory.stagingBuffer,488 (vk::Optional<const vk::AllocationCallbacks>)nullptr);489 }490 komputeManager()->device()->freeMemory(491 *memory.primaryMemory,492 (vk::Optional<const vk::AllocationCallbacks>)nullptr);493 if (memory.stagingMemory) {494 komputeManager()->device()->freeMemory(495 *memory.stagingMemory,496 (vk::Optional<const vk::AllocationCallbacks>)nullptr);497 }498}499 500static const char * ggml_backend_kompute_buffer_type_get_name(ggml_backend_buffer_type_t buft);501 502static503ggml_vk_memory * ggml_vk_find_tensor(const struct ggml_tensor * t, uint64_t & offset) {504 ggml_backend_buffer_t buffer = t->view_src ? t->view_src->buffer : t->buffer;505 506 // compatibility with ggml-backend507 GGML_ASSERT(buffer && buffer->buft->iface.get_name == ggml_backend_kompute_buffer_type_get_name);508 509 ggml_vk_memory * buf_ctx = static_cast<ggml_vk_memory *>(buffer->context);510 511 const intptr_t ioffs = intptr_t(t->data) - intptr_t(buf_ctx->data);512 513 GGML_ASSERT(ioffs >= 0 && ioffs + int64_t(ggml_nbytes(t)) <= int64_t(buffer->size));514 515 offset = uint64_t(ioffs);516 return buf_ctx;517}518 519static520const std::shared_ptr<kp::Tensor> ggml_vk_get_tensor(const struct ggml_tensor * t, uint32_t * alignedOffset = nullptr) {521 uint64_t originalOffset = 0;522 auto * res = ggml_vk_find_tensor(t, originalOffset);523 if (!res) {524 static std::shared_ptr<kp::Tensor> nullTensor = nullptr;525 return nullTensor;526 }527 528 // Create a tensor whose memory will be composed of our buffers at the correct offset529 const size_t nelements = ggml_nelements(t);530 size_t nbytes = ggml_nbytes(t);531 532 size_t vulkanOffset = ggml_vk_aligned_offset(t->buffer, originalOffset);533 if (alignedOffset) {534 *alignedOffset = originalOffset - vulkanOffset;535 nbytes += *alignedOffset;536 }537 538 return komputeManager()->tensor(539 t->data,540 nelements,541 nbytes, kp::Tensor::TensorDataTypes::eFloat,542 res->primaryMemory, res->primaryBuffer,543 res->stagingMemory, res->stagingBuffer,544 vulkanOffset);545}546 547static std::vector<uint32_t> getSpirvShader(const unsigned char* rawData, size_t size) {548 if (size % sizeof(uint32_t) != 0) {549 throw std::runtime_error("Invalid size: must be divisible by sizeof(uint32_t)");550 }551 552 const uint32_t* data_ptr = reinterpret_cast<const uint32_t*>(rawData);553 size_t count = size / sizeof(uint32_t);554 return std::vector<uint32_t>(data_ptr, data_ptr + count);555}556 557inline static558uint32_t safe_divide(uint32_t a, uint32_t b) {559 if (b <= 1) {560 return a;561 }562 if ((a % b) != 0) {563 fprintf(stderr, "((%u %% %u) == %u) != 0\n", a, b, a % b);564 GGML_ABORT("safe_divide result would've had remainder");565 }566 return a / b;567}568 569static void ggml_vk_add(570 kp::Sequence& seq,571 const std::shared_ptr<kp::Tensor>& inA,572 const std::shared_ptr<kp::Tensor>& inB,573 const std::shared_ptr<kp::Tensor>& out,574 uint32_t inAOff, uint32_t inBOff, uint32_t outOff,575 int32_t ne00, int32_t ne01, int32_t ne02, int32_t ne03,576 int32_t nb00, int32_t nb01, int32_t nb02, int32_t nb03,577 int32_t ne10, int32_t ne11, int32_t ne12, int32_t ne13,578 int32_t nb10, int32_t nb11, int32_t nb12, int32_t nb13,579 int32_t ne0,580 int32_t nb0, int32_t nb1, int32_t nb2, int32_t nb3581) {582 const static auto spirv = getSpirvShader(kp::shader_data::op_add_comp_spv,583 kp::shader_data::op_add_comp_spv_len);584 585 struct PushConstants {586 uint32_t inAOff, inBOff, outOff;587 int32_t ne00;588 int32_t nb00, nb01, nb02, nb03;589 int32_t ne10, ne11, ne12, ne13;590 int32_t nb10, nb11, nb12, nb13;591 int32_t ne0;592 int32_t nb0, nb1, nb2, nb3;593 } const pushConsts {594 safe_divide(inAOff, 4), safe_divide(inBOff, 4), safe_divide(outOff, 4),595 ne00,596 nb00, nb01, nb02, nb03,597 ne10, ne11, ne12, ne13,598 nb10, nb11, nb12, nb13,599 ne0,600 nb0, nb1, nb2, nb3601 };602 603 std::shared_ptr<kp::Algorithm> s_algo = nullptr;604 if (!komputeManager()->hasAlgorithm(__func__)) {605 s_algo = komputeManager()->algorithm<float, PushConstants>(__func__, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {unsigned(ne01), unsigned(ne02), unsigned(ne03)}, {}, {pushConsts});606 } else {607 s_algo = komputeManager()->getAlgorithm(__func__);608 s_algo->setTensors({inA, inB, out});609 s_algo->setWorkgroup({unsigned(ne01), unsigned(ne02), unsigned(ne03)});610 s_algo->setPushConstants<PushConstants>({pushConsts});611 s_algo->updateDescriptors(s_kompute_context->pool.get());612 }613 seq.record<kp::OpAlgoDispatch>(s_algo);614}615 616static void ggml_vk_addrow(kp::Sequence& seq,617 const std::shared_ptr<kp::Tensor>& inA,618 const std::shared_ptr<kp::Tensor>& inB,619 const std::shared_ptr<kp::Tensor>& out,620 uint32_t inAOff, uint32_t inBOff, uint32_t outOff,621 uint32_t size, uint32_t row = 0) {622 623 const static auto spirv = getSpirvShader(kp::shader_data::op_addrow_comp_spv,624 kp::shader_data::op_addrow_comp_spv_len);625 626 struct PushConstants {627 uint32_t inAOff, inBOff, outOff;628 uint32_t row;629 } const pushConsts {630 safe_divide(inAOff, 4), safe_divide(inBOff, 4), safe_divide(outOff, 4),631 row632 };633 634 std::shared_ptr<kp::Algorithm> s_algo = nullptr;635 if (!komputeManager()->hasAlgorithm(__func__))636 s_algo = komputeManager()->algorithm<float, PushConstants>(__func__, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {size}, {}, {pushConsts});637 else {638 s_algo = komputeManager()->getAlgorithm(__func__);639 s_algo->setTensors({inA, inB, out});640 s_algo->setWorkgroup({size});641 s_algo->setPushConstants<PushConstants>({pushConsts});642 s_algo->updateDescriptors(s_kompute_context->pool.get());643 }644 seq.record<kp::OpAlgoDispatch>(s_algo);645}646 647static void ggml_vk_mul(648 kp::Sequence& seq,649 const std::shared_ptr<kp::Tensor>& inA,650 const std::shared_ptr<kp::Tensor>& inB,651 const std::shared_ptr<kp::Tensor>& out,652 uint32_t inAOff, uint32_t inBOff, uint32_t outOff,653 int32_t ne00, int32_t ne01, int32_t ne02, int32_t ne03,654 int32_t nb00, int32_t nb01, int32_t nb02, int32_t nb03,655 int32_t ne10, int32_t ne11, int32_t ne12, int32_t ne13,656 int32_t nb10, int32_t nb11, int32_t nb12, int32_t nb13,657 int32_t ne0,658 int32_t nb0, int32_t nb1, int32_t nb2, int32_t nb3659) {660 const static auto spirv = getSpirvShader(kp::shader_data::op_mul_comp_spv,661 kp::shader_data::op_mul_comp_spv_len);662 663 struct PushConstants {664 uint32_t inAOff, inBOff, outOff;665 int32_t ne00;666 int32_t nb00, nb01, nb02, nb03;667 int32_t ne10, ne11, ne12, ne13;668 int32_t nb10, nb11, nb12, nb13;669 int32_t ne0;670 int32_t nb0, nb1, nb2, nb3;671 } const pushConsts {672 safe_divide(inAOff, 4), safe_divide(inBOff, 4), safe_divide(outOff, 4),673 ne00,674 nb00, nb01, nb02, nb03,675 ne10, ne11, ne12, ne13,676 nb10, nb11, nb12, nb13,677 ne0,678 nb0, nb1, nb2, nb3679 };680 681 std::shared_ptr<kp::Algorithm> s_algo = nullptr;682 if (!komputeManager()->hasAlgorithm(__func__)) {683 s_algo = komputeManager()->algorithm<float, PushConstants>(__func__, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {unsigned(ne01), unsigned(ne02), unsigned(ne03)}, {}, {pushConsts});684 } else {685 s_algo = komputeManager()->getAlgorithm(__func__);686 s_algo->setTensors({inA, inB, out});687 s_algo->setWorkgroup({unsigned(ne01), unsigned(ne02), unsigned(ne03)});688 s_algo->setPushConstants<PushConstants>({pushConsts});689 s_algo->updateDescriptors(s_kompute_context->pool.get());690 }691 seq.record<kp::OpAlgoDispatch>(s_algo);692}693 694static void ggml_vk_scale(kp::Sequence& seq,695 const std::shared_ptr<kp::Tensor>& in,696 const std::shared_ptr<kp::Tensor>& out,697 uint32_t inOff, uint32_t outOff,698 uint32_t size, float scale) {699 const static auto spirv_1 = getSpirvShader(700 kp::shader_data::op_scale_comp_spv, kp::shader_data::op_scale_comp_spv_len701 );702 const static auto spirv_8 = getSpirvShader(703 kp::shader_data::op_scale_8_comp_spv, kp::shader_data::op_scale_8_comp_spv_len704 );705 706 struct PushConstants {707 uint32_t inOff, outOff;708 float scale;709 } const pushConsts {710 safe_divide(inOff, 4), safe_divide(outOff, 4),711 scale712 };713 714 const auto * spirv = &spirv_1;715 std::string name(__func__);716 if (size % 8 == 0) {717 size /= 8;718 name += "_8";719 spirv = &spirv_8;720 }721 722 std::shared_ptr<kp::Algorithm> s_algo = nullptr;723 if (!komputeManager()->hasAlgorithm(name)) {724 s_algo = komputeManager()->algorithm<float, PushConstants>(name, s_kompute_context->pool.get(), {in, out}, *spirv, {size}, {}, {pushConsts});725 } else {726 s_algo = komputeManager()->getAlgorithm(name);727 s_algo->setTensors({in, out});728 s_algo->setWorkgroup({size});729 s_algo->setPushConstants<PushConstants>({pushConsts});730 s_algo->updateDescriptors(s_kompute_context->pool.get());731 }732 seq.record<kp::OpAlgoDispatch>(s_algo);733}734 735static void ggml_vk_xxlu(736 const std::vector<uint32_t>& spirv, const char * suffix, kp::Sequence& seq,737 const std::shared_ptr<kp::Tensor>& in,738 const std::shared_ptr<kp::Tensor>& out,739 uint32_t inOff, uint32_t outOff,740 uint32_t size741) {742 struct PushConstants {743 uint32_t inOff, outOff;744 } const pushConsts {745 safe_divide(inOff, 4), safe_divide(outOff, 4),746 };747 748 auto name = std::string(__func__) + "_" + suffix;749 std::shared_ptr<kp::Algorithm> s_algo = nullptr;750 if (!komputeManager()->hasAlgorithm(name)) {751 s_algo = komputeManager()->algorithm<float, PushConstants>(name, s_kompute_context->pool.get(), {in, out}, spirv, {size}, {}, {pushConsts});752 } else {753 s_algo = komputeManager()->getAlgorithm(name);754 s_algo->setTensors({in, out});755 s_algo->setWorkgroup({size});756 s_algo->setPushConstants<PushConstants>({pushConsts});757 s_algo->updateDescriptors(s_kompute_context->pool.get());758 }759 seq.record<kp::OpAlgoDispatch>(s_algo);760}761 762template <typename... Args>763static void ggml_vk_silu(Args&&... args) {764 const static auto spirv = getSpirvShader(kp::shader_data::op_silu_comp_spv,765 kp::shader_data::op_silu_comp_spv_len);766 767 ggml_vk_xxlu(spirv, "silu", std::forward<Args>(args)...);768}769 770template <typename... Args>771static void ggml_vk_relu(Args&&... args) {772 const static auto spirv = getSpirvShader(kp::shader_data::op_relu_comp_spv,773 kp::shader_data::op_relu_comp_spv_len);774 775 ggml_vk_xxlu(spirv, "relu", std::forward<Args>(args)...);776}777 778template <typename... Args>779static void ggml_vk_gelu(Args&&... args) {780 const static auto spirv = getSpirvShader(kp::shader_data::op_gelu_comp_spv,781 kp::shader_data::op_gelu_comp_spv_len);782 783 ggml_vk_xxlu(spirv, "gelu", std::forward<Args>(args)...);784}785 786static void ggml_vk_soft_max(787 kp::Sequence& seq,788 const std::shared_ptr<kp::Tensor>& inA,789 const std::shared_ptr<kp::Tensor>& inB,790 const std::shared_ptr<kp::Tensor>& out,791 uint32_t inAOff, uint32_t inBOff, uint32_t outOff,792 int32_t ne00, int32_t ne01, int32_t ne02, uint32_t ne03,793 float scale, float max_bias, float m0, float m1,794 uint32_t n_head_log2795) {796 const static auto spirv = getSpirvShader(kp::shader_data::op_softmax_comp_spv,797 kp::shader_data::op_softmax_comp_spv_len);798 799 struct PushConstants {800 uint32_t inAOff, inBOff, outOff;801 int32_t ne00, ne01, ne02;802 float scale, max_bias, m0, m1;803 uint32_t n_head_log2;804 int32_t mask;805 } pushConsts {806 safe_divide(inAOff, 4), safe_divide(inBOff, 4), safe_divide(outOff, 4),807 ne00, ne01, ne02,808 scale, max_bias, m0, m1,809 n_head_log2,810 bool(inB)811 };812 813 auto & inB_ = inB ? inB : inA;814 815 std::shared_ptr<kp::Algorithm> s_algo = nullptr;816 if (!komputeManager()->hasAlgorithm(__func__)) {817 // FIXME: The softmax kernel needs to be fixed to use the subgroupsize which can vary by device818 const uint32_t local_x = 32;819 s_algo = komputeManager()->algorithm<uint32_t, PushConstants>(__func__, s_kompute_context->pool.get(), {inA, inB_, out}, spirv, {unsigned(ne01), unsigned(ne02), unsigned(ne03)}, {local_x}, {pushConsts});820 } else {821 s_algo = komputeManager()->getAlgorithm(__func__);822 s_algo->setTensors({inA, inB_, out});823 s_algo->setWorkgroup({unsigned(ne01), unsigned(ne02), unsigned(ne03)});824 s_algo->setPushConstants<PushConstants>({pushConsts});825 s_algo->updateDescriptors(s_kompute_context->pool.get());826 }827 seq.record<kp::OpAlgoDispatch>(s_algo);828}829 830static void ggml_vk_norm_(831 const std::vector<uint32_t>& spirv, const char * suffix, kp::Sequence& seq,832 const std::shared_ptr<kp::Tensor>& in,833 const std::shared_ptr<kp::Tensor>& out,834 uint32_t inOff, uint32_t outOff,835 int32_t ne00, int32_t nb01,836 int32_t nrows, float epsilon837) {838 GGML_ASSERT(nb01%sizeof(float) == 0);839 GGML_ASSERT(ne00%sizeof(float) == 0);840 841 struct PushConstants {842 uint32_t inOff, outOff;843 uint32_t ne00, nb01;844 float eps;845 } pushConsts {846 safe_divide(inOff, 4), safe_divide(outOff, 4),847 (uint32_t)ne00, (uint32_t)nb01, epsilon848 };849 850 auto name = std::string(__func__) + "_" + suffix;851 std::shared_ptr<kp::Algorithm> s_algo = nullptr;852 if (!komputeManager()->hasAlgorithm(name)) {853 s_algo = komputeManager()->algorithm<float, PushConstants>(name, s_kompute_context->pool.get(), {in, out}, spirv, {(uint32_t)nrows}, {}, {pushConsts});854 } else {855 s_algo = komputeManager()->getAlgorithm(name);856 s_algo->setTensors({in, out});857 s_algo->setWorkgroup({(uint32_t)nrows});858 s_algo->setPushConstants<PushConstants>({pushConsts});859 s_algo->updateDescriptors(s_kompute_context->pool.get());860 }861 seq.record<kp::OpAlgoDispatch>(s_algo);862}863 864template <typename... Args>865static void ggml_vk_norm(Args&&... args) {866 const static auto spirv = getSpirvShader(kp::shader_data::op_norm_comp_spv,867 kp::shader_data::op_norm_comp_spv_len);868 869 ggml_vk_norm_(spirv, "norm", std::forward<Args>(args)...);870}871 872template <typename... Args>873static void ggml_vk_rms_norm(Args&&... args) {874 const static auto spirv = getSpirvShader(kp::shader_data::op_rmsnorm_comp_spv,875 kp::shader_data::op_rmsnorm_comp_spv_len);876 877 ggml_vk_norm_(spirv, "rms", std::forward<Args>(args)...);878}879 880static void ggml_vk_diag_mask_inf(kp::Sequence& seq,881 const std::shared_ptr<kp::Tensor>& in,882 const std::shared_ptr<kp::Tensor>& out,883 uint32_t inOff, uint32_t outOff,884 uint32_t n_past,885 int32_t ne00, int32_t ne01, int32_t ne02) {886 const static auto spirv = getSpirvShader(kp::shader_data::op_diagmask_comp_spv,887 kp::shader_data::op_diagmask_comp_spv_len);888 889 struct PushConstants {890 uint32_t inOff, outOff;891 uint32_t n_past;892 int32_t ne00, ne01;893 } pushConsts {894 safe_divide(inOff, 4), safe_divide(outOff, 4),895 n_past,896 ne00, ne01897 };898 899 std::shared_ptr<kp::Algorithm> s_algo = nullptr;900 if (!komputeManager()->hasAlgorithm(__func__))901 s_algo = komputeManager()->algorithm<float, PushConstants>(__func__, s_kompute_context->pool.get(), {in, out}, spirv, {unsigned(ne00), unsigned(ne01), unsigned(ne02)}, {}, {pushConsts});902 else {903 s_algo = komputeManager()->getAlgorithm(__func__);904 s_algo->setTensors({in, out});905 s_algo->setWorkgroup({unsigned(ne00), unsigned(ne01), unsigned(ne02)});906 s_algo->setPushConstants<PushConstants>({pushConsts});907 s_algo->updateDescriptors(s_kompute_context->pool.get());908 }909 seq.record<kp::OpAlgoDispatch>(s_algo);910}911 912static void ggml_vk_mul_mat_f16(913 kp::Sequence& seq,914 const std::shared_ptr<kp::Tensor>& inA,915 const std::shared_ptr<kp::Tensor>& inB,916 const std::shared_ptr<kp::Tensor>& out,917 uint32_t inAOff, uint32_t inBOff, uint32_t outOff,918 int32_t ne00, int32_t ne01, int32_t ne02,919 uint32_t nb00, uint32_t nb01, uint32_t nb02, uint32_t nb03,920 int32_t ne10, int32_t ne11, int32_t ne12, int32_t ne13,921 uint32_t nb10, uint32_t nb11, uint32_t nb12, uint32_t nb13,922 int32_t ne0, int32_t ne1,923 uint32_t r2, uint32_t r3924) {925 const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_f16_comp_spv,926 kp::shader_data::op_mul_mat_f16_comp_spv_len);927 928 struct PushConstants {929 uint32_t inAOff, inBOff, outOff;930 int32_t ne00, ne01, ne02;931 uint32_t nb00, nb01, nb02, nb03;932 int32_t ne10, ne11, ne12;933 uint32_t nb10, nb11, nb12, nb13;934 int32_t ne0, ne1;935 uint32_t r2, r3;936 } pushConsts {937 safe_divide(inAOff, 2), safe_divide(inBOff, 4), safe_divide(outOff, 4),938 ne00, ne01, ne02,939 nb00, nb01, nb02, nb03,940 ne10, ne11, ne12,941 nb10, nb11, nb12, nb13,942 ne0, ne1,943 r2, r3944 };945 946 const unsigned ny = unsigned((ne11 + 4 - 1)/4);947 948 std::shared_ptr<kp::Algorithm> s_algo = nullptr;949 if (!komputeManager()->hasAlgorithm(__func__)) {950 const uint32_t local_x = ggml_vk_current_device().subgroupSize * 2;951 s_algo = komputeManager()->algorithm<uint32_t, PushConstants>(__func__, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {unsigned(ne01), ny, unsigned(ne12*ne13)}, {local_x}, {pushConsts});952 } else {953 s_algo = komputeManager()->getAlgorithm(__func__);954 s_algo->setTensors({inA, inB, out});955 s_algo->setWorkgroup({unsigned(ne01), ny, unsigned(ne12*ne13)});956 s_algo->setPushConstants<PushConstants>({pushConsts});957 s_algo->updateDescriptors(s_kompute_context->pool.get());958 }959 seq.record<kp::OpAlgoDispatch>(s_algo);960}961 962static void ggml_vk_mul_mat_mat_f32(kp::Sequence& seq,963 const std::shared_ptr<kp::Tensor>& inA,964 const std::shared_ptr<kp::Tensor>& inB,965 const std::shared_ptr<kp::Tensor>& out,966 uint32_t inAOff, uint32_t inBOff, uint32_t outOff,967 int32_t ne00, int32_t ne01, int32_t ne02,968 uint32_t nb01, uint32_t nb02,969 int32_t ne11, int32_t ne12,970 uint32_t nb11, uint32_t nb12,971 uint32_t nb1, uint32_t nb2) {972 const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_mat_f32_comp_spv,973 kp::shader_data::op_mul_mat_mat_f32_comp_spv_len);974 975 struct PushConstants {976 uint32_t inAOff, inBOff, outOff;977 int32_t ne00, ne01, ne02, ne11, ne12;978 uint32_t nb01, nb02;979 uint32_t nb11, nb12;980 uint32_t nb1, nb2;981 } pushConsts {982 safe_divide(inAOff, 4), safe_divide(inBOff, 4), safe_divide(outOff, 4),983 ne00, ne01, ne02, ne11, ne12,984 nb01, nb02, nb11, nb12,985 nb1, nb2986 };987 988 const uint32_t local_x = ggml_vk_current_device().subgroupSize;989 std::shared_ptr<kp::Algorithm> s_algo = nullptr;990 if (!komputeManager()->hasAlgorithm(__func__)) {991 s_algo = komputeManager()->algorithm<uint32_t, PushConstants>(__func__, s_kompute_context->pool.get(),992 {inA, inB, out}, spirv,993 {unsigned(ne01),994 unsigned(ne11),995 unsigned(std::max(ne12, ne02))996 },997 {local_x},998 {pushConsts});999 } else {1000 s_algo = komputeManager()->getAlgorithm(__func__);1001 s_algo->setTensors({inA, inB, out});1002 s_algo->setWorkgroup({unsigned(ne01),1003 unsigned(ne11),1004 unsigned(std::max(ne12, ne02)),1005 });1006 s_algo->setPushConstants<PushConstants>({pushConsts});1007 s_algo->updateDescriptors(s_kompute_context->pool.get());1008 }1009 seq.record<kp::OpAlgoDispatch>(s_algo);1010}1011 1012static void ggml_vk_mul_mat_impl(1013 const std::vector<uint32_t>& spirv, const char * suffix, uint32_t block_size, kp::Sequence& seq,1014 const std::shared_ptr<kp::Tensor>& inA,1015 const std::shared_ptr<kp::Tensor>& inB,1016 const std::shared_ptr<kp::Tensor>& out,1017 uint32_t inAOff, uint32_t inBOff, uint32_t outOff,1018 int32_t ne00, int32_t ne01, int32_t ne02,1019 int32_t ne10, int32_t ne11, int32_t ne12, int32_t ne13,1020 int32_t ne0, int32_t ne1,1021 uint32_t nb01, uint32_t nb02, uint32_t nb03,1022 uint32_t nb11, uint32_t nb12, uint32_t nb13,1023 uint32_t r2, uint32_t r31024) {1025 struct PushConstants {1026 uint32_t inAOff, inBOff, outOff;1027 int32_t ne00, ne01, ne02;1028 int32_t ne10, ne12;1029 int32_t ne0, ne1;1030 uint32_t nb01, nb02, nb03;1031 uint32_t nb11, nb12, nb13;1032 uint32_t r2, r3;1033 } pushConsts {1034 safe_divide(inAOff, block_size), safe_divide(inBOff, 4), safe_divide(outOff, 4),1035 ne00, ne01, ne02,1036 ne10, ne12,1037 ne0, ne1,1038 nb01, nb02, nb03,1039 nb11, nb12, nb13,1040 r2, r31041 };1042 1043 auto name = std::string(__func__) + "_" + suffix;1044 std::shared_ptr<kp::Algorithm> s_algo = nullptr;1045 if (!komputeManager()->hasAlgorithm(name)) {1046 const uint32_t local_x = (ggml_vk_current_device().subgroupSize * 2) / 8;1047 s_algo = komputeManager()->algorithm<uint32_t, PushConstants>(name, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {unsigned((ne01 + 7)/8), unsigned(ne11), unsigned(ne12*ne13)}, {local_x}, {pushConsts});1048 } else {1049 s_algo = komputeManager()->getAlgorithm(name);1050 s_algo->setTensors({inA, inB, out});1051 s_algo->setWorkgroup({unsigned((ne01 + 7)/8), unsigned(ne11), unsigned(ne12*ne13)});1052 s_algo->setPushConstants<PushConstants>({pushConsts});1053 s_algo->updateDescriptors(s_kompute_context->pool.get());1054 }1055 seq.record<kp::OpAlgoDispatch>(s_algo);1056}1057 1058template <typename... Args>1059static void ggml_vk_mul_mat_q4_0(Args&&... args) {1060 const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q4_0_comp_spv,1061 kp::shader_data::op_mul_mat_q4_0_comp_spv_len);1062 1063 ggml_vk_mul_mat_impl(spirv, "q4_0", 1/*We access blocks unaligned*/, std::forward<Args>(args)...);1064}1065 1066template <typename... Args>1067static void ggml_vk_mul_mat_q4_1(Args&&... args) {1068 const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q4_1_comp_spv,1069 kp::shader_data::op_mul_mat_q4_1_comp_spv_len);1070 1071 ggml_vk_mul_mat_impl(spirv, "q4_1", 1/*We access blocks unaligned*/, std::forward<Args>(args)...);1072}1073 1074template <typename... Args>1075static void ggml_vk_mul_mat_q8_0(Args&&... args) {1076 const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q8_0_comp_spv,1077 kp::shader_data::op_mul_mat_q8_0_comp_spv_len);1078 1079 ggml_vk_mul_mat_impl(spirv, "q8_0", 1/*We access blocks unaligned*/, std::forward<Args>(args)...);1080}1081 1082static void ggml_vk_mul_mat_q4_k(1083 kp::Sequence& seq,1084 const std::shared_ptr<kp::Tensor>& inA,1085 const std::shared_ptr<kp::Tensor>& inB,1086 const std::shared_ptr<kp::Tensor>& out,1087 uint32_t inAOff, uint32_t inBOff, uint32_t outOff,1088 int32_t ne00, int32_t ne01, int32_t ne02,1089 int32_t ne10, int32_t ne11, int32_t ne12, int32_t ne13,1090 int32_t ne0, int32_t ne1,1091 uint32_t nb01, uint32_t nb02, uint32_t nb03,1092 uint32_t nb11, uint32_t nb12, uint32_t nb13,1093 uint32_t r2, uint32_t r31094) {1095 const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q4_k_comp_spv,1096 kp::shader_data::op_mul_mat_q4_k_comp_spv_len);1097 1098 struct PushConstants {1099 uint32_t inAOff, inBOff, outOff;1100 int32_t ne00, ne10, ne0, ne1, ne01, ne02, ne12;1101 uint32_t nb01, nb02, nb03, nb11, nb12, nb13;1102 uint32_t r2, r3;1103 } pushConsts {1104 inAOff, safe_divide(inBOff, 4), safe_divide(outOff, 4),1105 ne00, ne10, ne0, ne1, ne01, ne02, ne12,1106 nb01, nb02, nb03, nb11, nb12, nb13,1107 r2, r31108 };1109 1110 std::shared_ptr<kp::Algorithm> s_algo = nullptr;1111 if (!komputeManager()->hasAlgorithm(__func__)) {1112 s_algo = komputeManager()->algorithm<uint32_t, PushConstants>(__func__, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {unsigned((ne01 + 3)/4), unsigned(ne11), unsigned(ne12) * unsigned(ne13)}, {}, {pushConsts});1113 } else {1114 s_algo = komputeManager()->getAlgorithm(__func__);1115 s_algo->setTensors({inA, inB, out});1116 s_algo->setWorkgroup({unsigned((ne01 + 3)/4), unsigned(ne11), unsigned(ne12) * unsigned(ne13)});1117 s_algo->setPushConstants<PushConstants>({pushConsts});1118 s_algo->updateDescriptors(s_kompute_context->pool.get());1119 }1120 seq.record<kp::OpAlgoDispatch>(s_algo);1121}1122 1123static void ggml_vk_mul_mat_q6_k(1124 kp::Sequence& seq,1125 const std::shared_ptr<kp::Tensor>& inA,1126 const std::shared_ptr<kp::Tensor>& inB,1127 const std::shared_ptr<kp::Tensor>& out,1128 uint32_t inAOff, uint32_t inBOff, uint32_t outOff,1129 int32_t ne00, int32_t ne01, int32_t ne02,1130 int32_t ne10, int32_t ne11, int32_t ne12, int32_t ne13,1131 int32_t ne0, int32_t ne1,1132 uint32_t nb01, uint32_t nb02, uint32_t nb03,1133 uint32_t nb11, uint32_t nb12, uint32_t nb13,1134 uint32_t r2, uint32_t r31135) {1136 const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q6_k_comp_spv,1137 kp::shader_data::op_mul_mat_q6_k_comp_spv_len);1138 1139 struct PushConstants {1140 uint32_t inAOff, inBOff, outOff;1141 int32_t ne00, ne10, ne0, ne1, ne01, ne02, ne12;1142 uint32_t nb01, nb02, nb03, nb11, nb12, nb13;1143 uint32_t r2, r3;1144 } pushConsts {1145 inAOff, safe_divide(inBOff, 4), safe_divide(outOff, 4),1146 ne00, ne10, ne0, ne1, ne01, ne02, ne12,1147 nb01, nb02, nb03, nb11, nb12, nb13,1148 r2, r31149 };1150 1151 std::shared_ptr<kp::Algorithm> s_algo = nullptr;1152 if (!komputeManager()->hasAlgorithm(__func__)) {1153 const uint32_t local_x = 2;1154 const uint32_t local_y = ggml_vk_current_device().subgroupSize;1155 s_algo = komputeManager()->algorithm<uint32_t, PushConstants>(__func__, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {unsigned((ne01 + 1)/2), unsigned(ne11), unsigned(ne12)*unsigned(ne13)}, {local_x, local_y}, {pushConsts});1156 } else {1157 s_algo = komputeManager()->getAlgorithm(__func__);1158 s_algo->setTensors({inA, inB, out});1159 s_algo->setWorkgroup({unsigned((ne01 + 1)/2), unsigned(ne11), unsigned(ne12)*unsigned(ne13)});1160 s_algo->setPushConstants<PushConstants>({pushConsts});1161 s_algo->updateDescriptors(s_kompute_context->pool.get());1162 }1163 seq.record<kp::OpAlgoDispatch>(s_algo);1164}1165 1166static void ggml_vk_get_rows(1167 const std::vector<uint32_t>& spirv,1168 const char * suffix,1169 unsigned element_size, unsigned qk,1170 kp::Sequence& seq,1171 const std::shared_ptr<kp::Tensor>& inA,1172 const std::shared_ptr<kp::Tensor>& inB,1173 const std::shared_ptr<kp::Tensor>& out,1174 uint32_t inAOff, uint32_t inBOff, uint32_t outOff,1175 int32_t ne00, int32_t nb01, int32_t nb1,1176 uint32_t size1177) {1178 GGML_ASSERT(nb01%element_size == 0);1179 GGML_ASSERT(nb1%sizeof(float) == 0);1180 if (qk) GGML_ASSERT(ne00%qk == 0);1181 1182 struct PushConstants {1183 uint32_t inAOff, inBOff, outOff;1184 int32_t ne00, nb01, nb1;1185 } pushConsts {1186 safe_divide(inAOff, element_size), safe_divide(inBOff, 4), safe_divide(outOff, 4),1187 ne00, nb01, nb11188 };1189 1190 auto name = std::string(__func__) + "_" + suffix;1191 std::shared_ptr<kp::Algorithm> s_algo = nullptr;1192 if (!komputeManager()->hasAlgorithm(name)) {1193 s_algo = komputeManager()->algorithm<float, PushConstants>(name, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {size}, {}, {pushConsts});1194 } else {1195 s_algo = komputeManager()->getAlgorithm(name);1196 s_algo->setTensors({inA, inB, out});1197 s_algo->setWorkgroup({size});1198 s_algo->setPushConstants<PushConstants>({pushConsts});1199 s_algo->updateDescriptors(s_kompute_context->pool.get());1200 }