diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index ac021c7a6..c0ab9f1c6 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -2475,6 +2475,85 @@ static bool ggml_vk_strip_decode_vector(const uint32_t * code, size_t word_count return true; } +// Remove the loop unrolling hint of the matmul shader's BK loop +// and replace it with the dont_unroll hint for better performance on +// hardware like Apple M1/M2. +// Assumes 1. code comes from mul_mm.comp 2. the K-tile loop has no loop +// control hint and 3. the BK loop is the last loop nested directly inside +// the K-tile loop. +// Returns true when the input was modified; returns false otherwise +// without touching `out`. +static bool ggml_vk_roll_bk_loop(const uint32_t * code, size_t word_count, std::vector & out) { + if (word_count < 5) { + return false; + } + + struct vk_spv_loop { + size_t header; + size_t end; + uint32_t control; + }; + + std::vector loops; + + // Collect a list of all loops in the module. + for (size_t pos = 5; pos < word_count; ) { + const uint32_t wc = code[pos] >> spv::WordCountShift; + const uint32_t op = code[pos] & spv::OpCodeMask; + if (wc == 0 || pos + wc > word_count) { + return false; + } + + if (op == spv::OpLoopMerge && wc >= 4) { loops.push_back({ pos, 0, code[pos + 3] }); } + + if (op == spv::OpLabel && wc >= 2) { + for (auto & l : loops) { + if (l.end == 0 && code[l.header + 1] == code[pos + 1]) { l.end = pos; } + } + } + + pos += wc; + } + + auto encloses = [](const vk_spv_loop & a, const vk_spv_loop & b) { + return a.header < b.header && b.header < a.end; + }; + + // Find the BK loop. + const vk_spv_loop * bk = nullptr; + for (const auto & h : loops) { + if (h.control != spv::LoopControlUnrollMask) { + continue; + } + const vk_spv_loop * parent = nullptr; + bool has_child = false; + for (const auto & g : loops) { + if (encloses(g, h) && (!parent || g.header > parent->header)) { + parent = &g; + } + if (encloses(h, g)) { + has_child = true; + } + } + // BK loop should be the last loop nested inside the loop with no hint + // and have at least one child loop. + if (parent && + parent->control == spv::LoopControlMaskNone && + has_child && + (!bk || h.header > bk->header)) { + bk = &h; + } + } + if (!bk) { + return false; + } + + // set DontUnroll instead of Unroll + out.assign(code, code + word_count); + out[bk->header + 3] = spv::LoopControlDontUnrollMask; + return true; +} + static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipeline, size_t spv_size, const void* spv_data, const std::string entrypoint, uint32_t parameter_count, std::array wg_denoms, std::vector specialization_constants, bool disable_robustness, bool require_full_subgroups, uint32_t required_subgroup_size) { @@ -2558,6 +2637,22 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin } #endif +#if VK_HEADER_VERSION >= 287 + // Roll the mul_mm BK loop on Asahi Linux. Skip bf16 and the mul_mmq pipelines. + if (device->driver_id == vk::DriverId::eMesaHoneykrisp && + pipeline->name.rfind("matmul", 0) == 0 && + pipeline->name.find("bf16") == std::string::npos && + pipeline->name.find("q8_1") == std::string::npos) { + const uint32_t * src = spirv.empty() ? reinterpret_cast(spv_data) : spirv.data(); + size_t src_n = spirv.empty() ? spv_size / sizeof(uint32_t) : spirv.size(); + std::vector rolled; + if (ggml_vk_roll_bk_loop(src, src_n, rolled)) { + spirv = std::move(rolled); + shader_module_create_info = vk::ShaderModuleCreateInfo({}, spirv.size() * sizeof(uint32_t), spirv.data()); + } + } +#endif + pipeline->shader_module = device->device.createShaderModule(shader_module_create_info); vk::PushConstantRange pcr(