#include "ggml.h"
#include "ggml-backend.h"
#include "ggml-backend-impl.h"
#include "ggml-kompute.h"

// These are generated at build time by cmake custom command
#include "shaderop_scale.h"
#include "shaderop_scale_8.h"
#include "shaderop_add.h"
#include "shaderop_addrow.h"
#include "shaderop_mul.h"
#include "shaderop_silu.h"
#include "shaderop_relu.h"
#include "shaderop_gelu.h"
#include "shaderop_softmax.h"
#include "shaderop_norm.h"
#include "shaderop_rmsnorm.h"
#include "shaderop_diagmask.h"
#include "shaderop_mul_mat_f16.h"
#include "shaderop_mul_mat_q8_0.h"
#include "shaderop_mul_mat_q4_0.h"
#include "shaderop_mul_mat_q4_1.h"
#include "shaderop_mul_mat_q6_k.h"
#include "shaderop_mul_mat_mat_f32.h"
#include "shaderop_getrows_f16.h"
#include "shaderop_getrows_q4_0.h"
#include "shaderop_getrows_q4_1.h"
#include "shaderop_getrows_q6_k.h"
#include "shaderop_rope_f16.h"
#include "shaderop_rope_f32.h"
#include "shaderop_cpy_f16_f16.h"
#include "shaderop_cpy_f16_f32.h"
#include "shaderop_cpy_f32_f16.h"
#include "shaderop_cpy_f32_f32.h"

#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <list>
#include <memory>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include <kompute/Kompute.hpp>
#include <vulkan/vulkan.hpp>

#ifdef __linux__
#include <cstdlib> // for setenv
#endif

#define QK4_0 32
#define QR4_0 2
#define QK4_1 32
#define QK_NL 16

typedef ggml_fp16_t half;

static const std::shared_ptr<kp::Tensor> nullTensor = nullptr;

static std::string ggml_kompute_format_name(int device) {
    return "Kompute" + std::to_string(device);
}

struct ggml_kompute_context {
    int device;
    std::string name;
    std::shared_ptr<vk::DescriptorPool> pool;

    ggml_kompute_context(int device)
        : device(device), name(ggml_kompute_format_name(device)) {}
};

// FIXME: It would be good to consolidate the kompute manager and the kompute context into one object
// and consolidate the init functions and simplify object lifetime management. As it currently stands,
// we *have* to have the kompute manager no matter what for device discovery, but the kompute context
// is only created when a device is set and vulkan is explicitly turned on.
static ggml_kompute_context *s_kompute_context = nullptr;

class kompute_manager {
    kp::Manager *s_mgr = nullptr;

public:
    kp::Manager *operator()() {
        if (s_mgr && !s_mgr->hasInstance()) {
            destroy();
        }
        if (!s_mgr) {
            s_mgr = new kp::Manager;
        }
        return s_mgr;
    }

    void destroy() {
        delete s_mgr;
        s_mgr = nullptr;
    }
};

static kompute_manager komputeManager;

struct ggml_vk_memory {
    void *data = nullptr;
    size_t size = 0;
    vk::DeviceMemory *primaryMemory = nullptr;
    vk::Buffer *primaryBuffer = nullptr;
    vk::DeviceMemory *stagingMemory = nullptr;
    vk::Buffer *stagingBuffer = nullptr;
};

#ifdef __linux__
__attribute__((constructor))
static void enable_sam() {
    setenv("RADV_PERFTEST", "sam", false);
}
#endif

void ggml_vk_device_destroy(ggml_vk_device * device) {
    free(const_cast<char *>(device->name));
}

struct ggml_vk_device_cpp: ggml_vk_device {
    ggml_vk_device_cpp(
        int index, int type, size_t heapSize, const char * name, const char * vendor, int subgroupSize,
        uint64_t bufferAlignment, uint64_t maxAlloc
    )
        : ggml_vk_device({
              /* index           = */ index,
              /* type            = */ type,
              /* heapSize        = */ heapSize,
              /* name            = */ strdup(name),
              /* vendor          = */ vendor,
              /* supgroupSize    = */ subgroupSize,
              /* bufferAlignment = */ bufferAlignment,
              /* maxAlloc        = */ maxAlloc
        })
    {}

    ggml_vk_device_cpp(ggml_vk_device_cpp && other)
        : ggml_vk_device(other)
    {
        other.steal();
    }

    ggml_vk_device_cpp(ggml_vk_device_cpp &  other)
        : ggml_vk_device(other)
    {
        name = strdup(name);
    }

    ggml_vk_device_cpp & operator=(ggml_vk_device_cpp && other) {
        ggml_vk_device_destroy(this);
        static_cast<ggml_vk_device &>(*this) = other;
        other.steal();
        return *this;
    }

    ggml_vk_device_cpp & operator=(ggml_vk_device_cpp &  other) {
        ggml_vk_device_destroy(this);
        static_cast<ggml_vk_device &>(*this) = other;
        name = strdup(name);
        return *this;
    }

    ~ggml_vk_device_cpp() { ggml_vk_device_destroy(this); }

    // release the reference so we don't double-free name
    void steal() { name = nullptr; }
};

static bool ggml_vk_checkPhysicalDeviceFeatures(vk::PhysicalDevice physical_device) {
    vk::PhysicalDeviceFeatures availableFeatures;
    physical_device.getFeatures(&availableFeatures);

    if (!availableFeatures.shaderInt16)
        return false;

    vk::PhysicalDeviceVulkan11Features availableFeatures11;
    vk::PhysicalDeviceVulkan12Features availableFeatures12;

    availableFeatures11.pNext = &availableFeatures12;
    availableFeatures12.pNext = nullptr;

    vk::PhysicalDeviceFeatures2 features2;
    features2.pNext = &availableFeatures11;

    physical_device.getFeatures2(&features2);

    if (!availableFeatures11.uniformAndStorageBuffer16BitAccess ||
        !availableFeatures11.storageBuffer16BitAccess) {
        return false;
    }

    if (!availableFeatures12.storageBuffer8BitAccess ||
        !availableFeatures12.uniformAndStorageBuffer8BitAccess ||
        !availableFeatures12.shaderInt8) {
        return false;
    }

    return true;
}

static const char * ggml_vk_getVendorName(uint32_t vendorID) {
    switch (vendorID) {
        case 0x10DE:
            return "nvidia";
        case 0x1002:
            return "amd";
        case 0x8086:
            return "intel";
        default:
            return "unknown";
    }
}

static std::list<ggml_vk_device_cpp> ggml_vk_available_devices_internal(size_t memoryRequired) {
    std::list<ggml_vk_device_cpp> results;
    if (!komputeManager()->hasVulkan() || !komputeManager()->hasInstance())
        return results;

    std::vector<vk::PhysicalDevice> physical_devices;
    try {
        physical_devices = komputeManager()->listDevices();
    } catch (vk::SystemError & err) {
        std::cerr << __func__ << ": ignoring Vulkan exception: " << err.what() << "\n";
        return results;
    }

    uint32_t deviceCount = physical_devices.size();
    if (deviceCount == 0)
        return results;

    std::unordered_map<std::string, size_t> count_by_name;

    for (uint32_t i = 0; i < deviceCount; i++) {
        const auto & physical_device = physical_devices[i];

        VkPhysicalDeviceProperties dev_props = physical_device.getProperties();
        const uint32_t major = VK_VERSION_MAJOR(dev_props.apiVersion);
        const uint32_t minor = VK_VERSION_MINOR(dev_props.apiVersion);
        if (major < 1 || minor < 2)
            continue;

        if (dev_props.vendorID == 0x8086)
            continue; // Intel GPUs not supported

        if (!ggml_vk_checkPhysicalDeviceFeatures(physical_device))
            continue;

        VkPhysicalDeviceMemoryProperties memoryProperties = physical_device.getMemoryProperties();

        size_t heapSize = 0;
        for (uint32_t j = 0; j < memoryProperties.memoryHeapCount; ++j) {
            VkMemoryHeap heap = memoryProperties.memoryHeaps[j];
            if (heap.flags & VK_MEMORY_HEAP_DEVICE_LOCAL_BIT) {
                heapSize = heap.size;
                break;
            }
        }

        if (heapSize < memoryRequired)
            continue;

        auto ext_props = physical_device.enumerateDeviceExtensionProperties();
        bool has_maintenance4 = false;

        // Check if maintenance4 is supported
        for (const auto & properties : ext_props) {
            if (strcmp("VK_KHR_maintenance4", properties.extensionName) == 0) {
                has_maintenance4 = true;
            }
        }

        vk::PhysicalDeviceSubgroupProperties subgroup_props;
        vk::PhysicalDeviceProperties2 dev_props2;
        vk::PhysicalDeviceMaintenance3Properties dev_props3;
        vk::PhysicalDeviceMaintenance4Properties dev_props4;
        dev_props2.pNext = &dev_props3;
        dev_props3.pNext = &subgroup_props;
        if (has_maintenance4) {
            subgroup_props.pNext = &dev_props4;
        }
        physical_device.getProperties2(&dev_props2);

        if (subgroup_props.subgroupSize < 32)
            continue;

        std::string name(dev_props.deviceName);
        size_t n_idx = ++count_by_name[name];
        if (n_idx > 1) {
            name += " (" + std::to_string(n_idx) + ")";
        }

        uint64_t maxAlloc = dev_props3.maxMemoryAllocationSize;
        if (has_maintenance4) {
            maxAlloc = std::min(maxAlloc, dev_props4.maxBufferSize);
        }

        results.emplace_back(
            /* index           = */ i,
            /* type            = */ dev_props.deviceType,
            /* heapSize        = */ heapSize,
            /* name            = */ name.c_str(),
            /* vendor          = */ ggml_vk_getVendorName(dev_props.vendorID),
            /* subgroupSize    = */ subgroup_props.subgroupSize,
            /* bufferAlignment = */ dev_props.limits.minStorageBufferOffsetAlignment,
            /* maxAlloc        = */ maxAlloc
        );
    }

    // std::list::sort is guaranteed to be stable
    results.sort(
        [](const ggml_vk_device_cpp & lhs, const ggml_vk_device_cpp & rhs) -> bool {
            if (lhs.type != rhs.type) {
                if (lhs.type == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) return true;
                if (rhs.type == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) return false;

                if (lhs.type == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU) return true;
                if (rhs.type == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU) return false;
            }
            return lhs.heapSize < rhs.heapSize;
        }
    );

    return results;
}

// public API returns a C-style array
ggml_vk_device * ggml_vk_available_devices(size_t memoryRequired, size_t * count) {
    auto devices = ggml_vk_available_devices_internal(memoryRequired);
    *count = devices.size();
    if (devices.empty()) {
        return nullptr;
    }

    size_t nbytes = sizeof (ggml_vk_device) * (devices.size());
    auto * arr = static_cast<ggml_vk_device *>(malloc(nbytes));

    int i = 0;
    for (auto & d : devices) { arr[i++] = d; d.steal(); }

    return arr;
}

static void ggml_vk_filterByVendor(std::list<ggml_vk_device_cpp> & devices, const std::string & targetVendor) {
    devices.erase(
        std::remove_if(devices.begin(), devices.end(),
            [&targetVendor](const ggml_vk_device_cpp & device) {
                return device.vendor != targetVendor;
            }),
        devices.end()
    );
}

static void ggml_vk_filterByName(std::list<ggml_vk_device_cpp> & devices, const std::string & targetName) {
    devices.erase(
        std::remove_if(devices.begin(), devices.end(),
            [&targetName](const ggml_vk_device_cpp & device) {
                return device.name != targetName;
            }),
        devices.end()
    );
}

bool ggml_vk_get_device(ggml_vk_device * device, size_t memoryRequired, const char * name) {
    if (!*name)
        return false;

    auto devices = ggml_vk_available_devices_internal(memoryRequired);
    std::string name_str(name);
    if (name_str == "amd" || name_str == "nvidia" || name_str == "intel") {
        ggml_vk_filterByVendor(devices, name_str);
    } else if (name_str != "gpu") {
        ggml_vk_filterByName(devices, name_str);
    }

    if (devices.empty())
        return false;

    *device = devices.front();
    return true;
}

bool ggml_vk_has_vulkan() {
    return komputeManager()->hasVulkan();
}

bool ggml_vk_has_device() {
    return komputeManager()->hasDevice();
}

ggml_vk_device ggml_vk_current_device() {
    if (!komputeManager()->hasDevice())
        return ggml_vk_device();

    auto devices = ggml_vk_available_devices_internal(0);
    ggml_vk_filterByName(devices, komputeManager()->physicalDevice()->getProperties().deviceName.data());
    GGML_ASSERT(!devices.empty());

    ggml_vk_device device = devices.front();
    devices.front().steal();
    return device;
}

static
void ggml_vk_allocate_descriptor_pool(struct ggml_kompute_context * ctx, size_t size) {
    std::vector<vk::DescriptorPoolSize> descriptorPoolSizes = {
        vk::DescriptorPoolSize(
          vk::DescriptorType::eStorageBuffer,
          3 * size // Descriptor count is number of possible tensors to pass into an algorithm
          )
    };

    vk::DescriptorPoolCreateInfo descriptorPoolInfo(
      vk::DescriptorPoolCreateFlags(),
      size, // Max sets
      static_cast<uint32_t>(descriptorPoolSizes.size()),
      descriptorPoolSizes.data());

    ctx->pool = std::make_shared<vk::DescriptorPool>();
    vk::Result r = komputeManager()->device()->createDescriptorPool(
      &descriptorPoolInfo, nullptr, ctx->pool.get());
    if (r != vk::Result::eSuccess)
        std::cerr << "Error allocating descriptor pool" << vk::to_string(r);
}

static
void ggml_vk_free_descriptor_pool(struct ggml_kompute_context * ctx) {
    if (ctx->pool) {
        komputeManager()->device()->destroy(
          *ctx->pool,
          (vk::Optional<const vk::AllocationCallbacks>)nullptr);
        ctx->pool = nullptr;
    }
}

static
vk::Buffer *ggml_vk_allocate_buffer(size_t size) {
    vk::BufferCreateInfo bufferCreateInfo;
    bufferCreateInfo.size = size;
    bufferCreateInfo.usage = vk::BufferUsageFlagBits::eStorageBuffer |
                             vk::BufferUsageFlagBits::eTransferSrc |
                             vk::BufferUsageFlagBits::eTransferDst;
    bufferCreateInfo.sharingMode = vk::SharingMode::eExclusive;

    vk::Buffer *vkBuffer = new vk::Buffer;
    vk::Result r = komputeManager()->device()->createBuffer(&bufferCreateInfo, nullptr, vkBuffer);
    if (r != vk::Result::eSuccess)
        std::cerr << "Error allocating buffer " << vk::to_string(r) << std::endl;
    return vkBuffer;
}

static
vk::DeviceMemory *ggml_vk_allocate(size_t size, vk::MemoryPropertyFlags flags, vk::MemoryRequirements requirements, bool *isHostVisible) {

    uint32_t memoryTypeIndex = -1;
    bool memoryTypeIndexFound = false;
    vk::PhysicalDeviceMemoryProperties memoryProperties = komputeManager()->physicalDevice()->getMemoryProperties();
    for (uint32_t i = 0; i < memoryProperties.memoryTypeCount; i++) {
        const vk::MemoryType &memoryType = memoryProperties.memoryTypes[i];
        const vk::MemoryHeap &memoryHeap = memoryProperties.memoryHeaps[memoryType.heapIndex];
        if (memoryHeap.size < size) {
            continue;
        }

        if (requirements.memoryTypeBits & (1 << i)) {
            if (((memoryProperties.memoryTypes[i]).propertyFlags &
                 flags) == flags) {
                memoryTypeIndex = i;
                memoryTypeIndexFound = true;
                if (isHostVisible && (memoryProperties.memoryTypes[i].propertyFlags & vk::MemoryPropertyFlagBits::eHostVisible)) {
                    *isHostVisible = true;
                }
                break;
            }
        }
    }
    if (!memoryTypeIndexFound) {
        throw std::runtime_error(
          "Memory type index for buffer creation not found");
    }

    vk::MemoryAllocateInfo allocInfo;
    allocInfo.allocationSize = size;
    allocInfo.memoryTypeIndex = memoryTypeIndex;
    vk::DeviceMemory *vkDeviceMemory =  new vk::DeviceMemory;
    vk::Result r = komputeManager()->device()->allocateMemory(&allocInfo, nullptr, vkDeviceMemory);
    if (r != vk::Result::eSuccess) {
        std::cerr << "Error allocating memory " << vk::to_string(r) << std::endl;
        return nullptr;
    }
    return vkDeviceMemory;
}

static size_t ggml_vk_aligned_offset(ggml_backend_buffer_t buffer, size_t offset) {
    size_t minStorageBufferOffsetAlignment = ggml_backend_buffer_get_alignment(buffer);

    // If offset is already aligned, return it directly
    if (offset % minStorageBufferOffsetAlignment == 0) {
        return offset;
    }

    // Otherwise, return the largest multiple of minStorageBufferOffsetAlignment less than offset
    return (offset / minStorageBufferOffsetAlignment) * minStorageBufferOffsetAlignment;
}

static ggml_vk_memory ggml_vk_allocate(size_t size) {
    ggml_vk_memory memory;
    bool isHostVisible = false;
    {
        memory.primaryBuffer = ggml_vk_allocate_buffer(size);

        vk::MemoryRequirements memoryRequirements = komputeManager()->device()->getBufferMemoryRequirements(*memory.primaryBuffer);
        vk::MemoryPropertyFlags memoryPropertyFlags = vk::MemoryPropertyFlagBits::eDeviceLocal;
        memory.primaryMemory = ggml_vk_allocate(size, memoryPropertyFlags, memoryRequirements, &isHostVisible);
        if (!memory.primaryMemory)
            return {};

        komputeManager()->device()->bindBufferMemory(*memory.primaryBuffer, *memory.primaryMemory, 0);
        if (isHostVisible) {
            vk::Result r = komputeManager()->device()->mapMemory(*memory.primaryMemory, 0, size, vk::MemoryMapFlags(), &memory.data);
            if (r != vk::Result::eSuccess)
                std::cerr << "Error mapping memory" << vk::to_string(r);
        }
    }

    if (!isHostVisible) {
        memory.stagingBuffer = ggml_vk_allocate_buffer(size);
        vk::MemoryRequirements memoryRequirements = komputeManager()->device()->getBufferMemoryRequirements(*memory.stagingBuffer);
        vk::MemoryPropertyFlags memoryPropertyFlags = vk::MemoryPropertyFlagBits::eHostVisible |
                                                      vk::MemoryPropertyFlagBits::eHostCoherent |
                                                      vk::MemoryPropertyFlagBits::eHostCached;
        memory.stagingMemory = ggml_vk_allocate(size, memoryPropertyFlags, memoryRequirements, &isHostVisible);
        komputeManager()->device()->bindBufferMemory(*memory.stagingBuffer, *memory.stagingMemory, 0);
        vk::Result r = komputeManager()->device()->mapMemory(*memory.stagingMemory, 0, size, vk::MemoryMapFlags(), &memory.data);
        if (r != vk::Result::eSuccess)
            std::cerr << "Error mapping memory" << vk::to_string(r);
    }

    memory.size = size;
    return memory;
}

static void ggml_vk_free_memory(ggml_vk_memory &memory)
{
    komputeManager()->device()->destroy(
      *memory.primaryBuffer,
      (vk::Optional<const vk::AllocationCallbacks>)nullptr);
    if (memory.stagingBuffer) {
        komputeManager()->device()->destroy(
          *memory.stagingBuffer,
          (vk::Optional<const vk::AllocationCallbacks>)nullptr);
    }
    komputeManager()->device()->freeMemory(
      *memory.primaryMemory,
      (vk::Optional<const vk::AllocationCallbacks>)nullptr);
    if (memory.stagingMemory) {
        komputeManager()->device()->freeMemory(
          *memory.stagingMemory,
          (vk::Optional<const vk::AllocationCallbacks>)nullptr);
    }
}

static const char * ggml_backend_kompute_buffer_type_get_name(ggml_backend_buffer_type_t buft);

static
ggml_vk_memory * ggml_vk_find_tensor(const struct ggml_tensor * t, uint64_t & offset) {
    ggml_backend_buffer_t buffer = t->view_src ? t->view_src->buffer : t->buffer;

    // compatibility with ggml-backend
    GGML_ASSERT(buffer && buffer->buft->iface.get_name == ggml_backend_kompute_buffer_type_get_name);

    ggml_vk_memory * buf_ctx = static_cast<ggml_vk_memory *>(buffer->context);

    const intptr_t ioffs = intptr_t(t->data) - intptr_t(buf_ctx->data);

    GGML_ASSERT(ioffs >= 0 && ioffs + int64_t(ggml_nbytes(t)) <= int64_t(buffer->size));

    offset = uint64_t(ioffs);
    return buf_ctx;
}

static
const std::shared_ptr<kp::Tensor> ggml_vk_get_tensor_aligned(const struct ggml_tensor * t, uint32_t * aligned_offset) {
    uint64_t original_offset = 0;
    auto * res = ggml_vk_find_tensor(t, original_offset);
    if (!res) {
        return nullTensor;
    }

    // Create a tensor whose memory will be composed of our buffers at the correct offset
    size_t nbytes = ggml_nbytes(t);
    size_t vulkan_offset = ggml_vk_aligned_offset(t->buffer, original_offset);
    *aligned_offset = original_offset - vulkan_offset;
    nbytes += *aligned_offset;

    return komputeManager()->tensor(
        t->data,
        ggml_nelements(t), nbytes,
        kp::Tensor::TensorDataTypes::eFloat,
        res->primaryMemory, res->primaryBuffer,
        res->stagingMemory, res->stagingBuffer,
        vulkan_offset);
}

static
const std::shared_ptr<kp::Tensor> ggml_vk_get_tensor_slice(const struct ggml_tensor * t, size_t offset, size_t nbytes) {
    uint64_t tensor_offset = 0;
    auto * res = ggml_vk_find_tensor(t, tensor_offset);
    if (!res) {
        return nullTensor;
    }

    size_t elsz = ggml_element_size(t);
    GGML_ASSERT(nbytes % elsz == 0);

    // Create a tensor whose memory will be composed of our buffers at the correct offset
    return komputeManager()->tensor(
        reinterpret_cast<char *>(t->data) + offset,
        nbytes / elsz, nbytes,
        kp::Tensor::TensorDataTypes::eFloat,
        res->primaryMemory, res->primaryBuffer,
        res->stagingMemory, res->stagingBuffer,
        tensor_offset + offset);
}

static std::vector<uint32_t> getSpirvShader(const unsigned char* rawData, size_t size) {
    if (size % sizeof(uint32_t) != 0) {
        throw std::runtime_error("Invalid size: must be divisible by sizeof(uint32_t)");
    }

    const uint32_t* data_ptr = reinterpret_cast<const uint32_t*>(rawData);
    size_t count = size / sizeof(uint32_t);
    return std::vector<uint32_t>(data_ptr, data_ptr + count);
}

inline static
uint32_t safe_divide(uint32_t a, uint32_t b) {
    if (b <= 1) {
        return a;
    }
    if ((a % b) != 0) {
        fprintf(stderr, "((%u %% %u) == %u) != 0\n", a, b, a % b);
        GGML_ASSERT(!"safe_divide result would've had remainder");
    }
    return a / b;
}

static void ggml_vk_add(
    kp::Sequence& seq,
    const std::shared_ptr<kp::Tensor>& inA,
    const std::shared_ptr<kp::Tensor>& inB,
    const std::shared_ptr<kp::Tensor>& out,
    uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
    int32_t ne00, int32_t ne01, int32_t ne02, int32_t ne03,
    int32_t nb00, int32_t nb01, int32_t nb02, int32_t nb03,
    int32_t ne10, int32_t ne11, int32_t ne12, int32_t ne13,
    int32_t nb10, int32_t nb11, int32_t nb12, int32_t nb13,
    int32_t ne0,
    int32_t nb0,  int32_t nb1,  int32_t nb2,  int32_t nb3
) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_add_comp_spv,
        kp::shader_data::op_add_comp_spv_len);

    struct PushConstants {
        uint32_t inAOff, inBOff, outOff;
        int32_t ne00;
        int32_t nb00, nb01, nb02, nb03;
        int32_t ne10, ne11, ne12, ne13;
        int32_t nb10, nb11, nb12, nb13;
        int32_t ne0;
        int32_t nb0, nb1, nb2, nb3;
    } const pushConsts {
        safe_divide(inAOff, 4), safe_divide(inBOff, 4), safe_divide(outOff, 4),
        ne00,
        nb00, nb01, nb02, nb03,
        ne10, ne11, ne12, ne13,
        nb10, nb11, nb12, nb13,
        ne0,
        nb0, nb1, nb2, nb3
    };

    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(__func__)) {
        s_algo = komputeManager()->algorithm<float, PushConstants>(__func__, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {unsigned(ne01), unsigned(ne02), unsigned(ne03)}, {}, {pushConsts});
    } else {
        s_algo = komputeManager()->getAlgorithm(__func__);
        s_algo->setTensors({inA, inB, out});
        s_algo->setWorkgroup({unsigned(ne01), unsigned(ne02), unsigned(ne03)});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

static void ggml_vk_addrow(kp::Sequence& seq,
                 const std::shared_ptr<kp::Tensor>& inA,
                 const std::shared_ptr<kp::Tensor>& inB,
                 const std::shared_ptr<kp::Tensor>& out,
                 uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
                 uint32_t size, uint32_t row = 0) {

    const static auto spirv = getSpirvShader(kp::shader_data::op_addrow_comp_spv,
        kp::shader_data::op_addrow_comp_spv_len);

    struct PushConstants {
        uint32_t inAOff, inBOff, outOff;
        uint32_t row;
    } const pushConsts {
        safe_divide(inAOff, 4), safe_divide(inBOff, 4), safe_divide(outOff, 4),
        row
    };

    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(__func__))
        s_algo = komputeManager()->algorithm<float, PushConstants>(__func__, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {size}, {}, {pushConsts});
    else {
        s_algo = komputeManager()->getAlgorithm(__func__);
        s_algo->setTensors({inA, inB, out});
        s_algo->setWorkgroup({size});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

static void ggml_vk_mul(
    kp::Sequence& seq,
    const std::shared_ptr<kp::Tensor>& inA,
    const std::shared_ptr<kp::Tensor>& inB,
    const std::shared_ptr<kp::Tensor>& out,
    uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
    int32_t ne00, int32_t ne01, int32_t ne02, int32_t ne03,
    int32_t nb00, int32_t nb01, int32_t nb02, int32_t nb03,
    int32_t ne10, int32_t ne11, int32_t ne12, int32_t ne13,
    int32_t nb10, int32_t nb11, int32_t nb12, int32_t nb13,
    int32_t ne0,
    int32_t nb0,  int32_t nb1,  int32_t nb2,  int32_t nb3
) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_mul_comp_spv,
        kp::shader_data::op_mul_comp_spv_len);

    struct PushConstants {
        uint32_t inAOff, inBOff, outOff;
        int32_t ne00;
        int32_t nb00, nb01, nb02, nb03;
        int32_t ne10, ne11, ne12, ne13;
        int32_t nb10, nb11, nb12, nb13;
        int32_t ne0;
        int32_t nb0, nb1, nb2, nb3;
    } const pushConsts {
        safe_divide(inAOff, 4), safe_divide(inBOff, 4), safe_divide(outOff, 4),
        ne00,
        nb00, nb01, nb02, nb03,
        ne10, ne11, ne12, ne13,
        nb10, nb11, nb12, nb13,
        ne0,
        nb0, nb1, nb2, nb3
    };

    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(__func__)) {
        s_algo = komputeManager()->algorithm<float, PushConstants>(__func__, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {unsigned(ne01), unsigned(ne02), unsigned(ne03)}, {}, {pushConsts});
    } else {
        s_algo = komputeManager()->getAlgorithm(__func__);
        s_algo->setTensors({inA, inB, out});
        s_algo->setWorkgroup({unsigned(ne01), unsigned(ne02), unsigned(ne03)});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

static void ggml_vk_scale(kp::Sequence& seq,
                   const std::shared_ptr<kp::Tensor>& in,
                   const std::shared_ptr<kp::Tensor>& out,
                   uint32_t inOff, uint32_t outOff,
                   uint32_t size, float scale) {
    const static auto spirv_1 = getSpirvShader(
        kp::shader_data::op_scale_comp_spv, kp::shader_data::op_scale_comp_spv_len
    );
    const static auto spirv_8 = getSpirvShader(
        kp::shader_data::op_scale_8_comp_spv, kp::shader_data::op_scale_8_comp_spv_len
    );

    struct PushConstants {
        uint32_t inOff, outOff;
        float scale;
    } const pushConsts {
        safe_divide(inOff, 4), safe_divide(outOff, 4),
        scale
    };

    const auto * spirv = &spirv_1;
    std::string name(__func__);
    if (size % 8 == 0) {
        size /= 8;
        name += "_8";
        spirv = &spirv_8;
    }

    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(name)) {
        s_algo = komputeManager()->algorithm<float, PushConstants>(name, s_kompute_context->pool.get(), {in, out}, *spirv, {size}, {}, {pushConsts});
    } else {
        s_algo = komputeManager()->getAlgorithm(name);
        s_algo->setTensors({in, out});
        s_algo->setWorkgroup({size});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

static void ggml_vk_xxlu(
    const std::vector<uint32_t>& spirv, const char * suffix, kp::Sequence& seq,
    const std::shared_ptr<kp::Tensor>& in,
    const std::shared_ptr<kp::Tensor>& out,
    uint32_t inOff, uint32_t outOff,
    uint32_t size
) {
    struct PushConstants {
        uint32_t inOff, outOff;
    } const pushConsts {
        safe_divide(inOff, 4), safe_divide(outOff, 4),
    };

    auto name = std::string(__func__) + "_" + suffix;
    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(name)) {
        s_algo = komputeManager()->algorithm<float, PushConstants>(name, s_kompute_context->pool.get(), {in, out}, spirv, {size}, {}, {pushConsts});
    } else {
        s_algo = komputeManager()->getAlgorithm(name);
        s_algo->setTensors({in, out});
        s_algo->setWorkgroup({size});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

template <typename... Args>
static void ggml_vk_silu(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_silu_comp_spv,
        kp::shader_data::op_silu_comp_spv_len);

    ggml_vk_xxlu(spirv, "silu", std::forward<Args>(args)...);
}

template <typename... Args>
static void ggml_vk_relu(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_relu_comp_spv,
        kp::shader_data::op_relu_comp_spv_len);

    ggml_vk_xxlu(spirv, "relu", std::forward<Args>(args)...);
}

template <typename... Args>
static void ggml_vk_gelu(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_gelu_comp_spv,
        kp::shader_data::op_gelu_comp_spv_len);

    ggml_vk_xxlu(spirv, "gelu", std::forward<Args>(args)...);
}

static void ggml_vk_soft_max(
    kp::Sequence& seq,
    const std::shared_ptr<kp::Tensor>& inA,
    const std::shared_ptr<kp::Tensor>& inB,
    const std::shared_ptr<kp::Tensor>& out,
    uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
    int32_t ne00, int32_t ne01, int32_t ne02, uint32_t ne03,
    float scale
) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_softmax_comp_spv,
        kp::shader_data::op_softmax_comp_spv_len);

    struct PushConstants {
        uint32_t inAOff, inBOff, outOff;
        int32_t ne00, ne01, ne02;
        float scale;
        int32_t mask;
    } pushConsts {
        safe_divide(inAOff, 4), safe_divide(inBOff, 4), safe_divide(outOff, 4),
        ne00, ne01, ne02,
        scale,
        bool(inB)
    };

    auto & inB_ = inB ? inB : inA;

    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(__func__)) {
        // FIXME: The softmax kernel needs to be fixed to use the subgroupsize which can vary by device
        const uint32_t local_x = 32;
        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});
    } else {
        s_algo = komputeManager()->getAlgorithm(__func__);
        s_algo->setTensors({inA, inB_, out});
        s_algo->setWorkgroup({unsigned(ne01), unsigned(ne02), unsigned(ne03)});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

static void ggml_vk_norm_(
    const std::vector<uint32_t>& spirv, const char * suffix, kp::Sequence& seq,
    const std::shared_ptr<kp::Tensor>& in,
    const std::shared_ptr<kp::Tensor>& out,
    uint32_t inOff, uint32_t outOff,
    int32_t ne00, int32_t nb01,
    int32_t nrows, float epsilon
) {
    GGML_ASSERT(nb01%sizeof(float) == 0);
    GGML_ASSERT(ne00%sizeof(float) == 0);

    struct PushConstants {
        uint32_t inOff, outOff;
        uint32_t ne00, nb01;
        float eps;
    } pushConsts {
        safe_divide(inOff, 4), safe_divide(outOff, 4),
        (uint32_t)ne00, (uint32_t)nb01, epsilon
    };

    auto name = std::string(__func__) + "_" + suffix;
    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(name)) {
        s_algo = komputeManager()->algorithm<float, PushConstants>(name, s_kompute_context->pool.get(), {in, out}, spirv, {(uint32_t)nrows}, {}, {pushConsts});
    } else {
        s_algo = komputeManager()->getAlgorithm(name);
        s_algo->setTensors({in, out});
        s_algo->setWorkgroup({(uint32_t)nrows});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

template <typename... Args>
static void ggml_vk_norm(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_norm_comp_spv,
        kp::shader_data::op_norm_comp_spv_len);

    ggml_vk_norm_(spirv, "norm", std::forward<Args>(args)...);
}

template <typename... Args>
static void ggml_vk_rms_norm(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_rmsnorm_comp_spv,
        kp::shader_data::op_rmsnorm_comp_spv_len);

    ggml_vk_norm_(spirv, "rms", std::forward<Args>(args)...);
}

static void ggml_vk_diag_mask_inf(kp::Sequence& seq,
                           const std::shared_ptr<kp::Tensor>& in,
                           const std::shared_ptr<kp::Tensor>& out,
                           uint32_t inOff, uint32_t outOff,
                           uint32_t n_past,
                           int32_t ne00, int32_t ne01, int32_t ne02) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_diagmask_comp_spv,
        kp::shader_data::op_diagmask_comp_spv_len);

    struct PushConstants {
        uint32_t inOff, outOff;
        uint32_t n_past;
        int32_t ne00, ne01;
    } pushConsts {
        safe_divide(inOff, 4), safe_divide(outOff, 4),
        n_past,
        ne00, ne01
    };

    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(__func__))
        s_algo = komputeManager()->algorithm<float, PushConstants>(__func__, s_kompute_context->pool.get(), {in, out}, spirv, {unsigned(ne00), unsigned(ne01), unsigned(ne02)}, {}, {pushConsts});
    else {
        s_algo = komputeManager()->getAlgorithm(__func__);
        s_algo->setTensors({in, out});
        s_algo->setWorkgroup({unsigned(ne00), unsigned(ne01), unsigned(ne02)});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

static void ggml_vk_mul_mat_f16(
    kp::Sequence& seq,
    const std::shared_ptr<kp::Tensor>& inA,
    const std::shared_ptr<kp::Tensor>& inB,
    const std::shared_ptr<kp::Tensor>& out,
    uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
    int32_t ne00, int32_t ne01, int32_t ne02,
    uint32_t nb00, uint32_t nb01, uint32_t nb02,
    int32_t ne10, int32_t ne11, int32_t ne12, int32_t ne13,
    uint32_t nb10, uint32_t nb11, uint32_t nb12,
    int32_t ne0, int32_t ne1,
    uint32_t r2, uint32_t r3
) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_f16_comp_spv,
        kp::shader_data::op_mul_mat_f16_comp_spv_len);

    struct PushConstants {
        uint32_t inAOff, inBOff, outOff;
        int32_t ne00, ne01, ne02;
        uint32_t nb00, nb01, nb02;
        int32_t ne10, ne11, ne12;
        uint32_t nb10, nb11, nb12;
        int32_t ne0, ne1;
        uint32_t r2, r3;
    } pushConsts {
        safe_divide(inAOff, 2), safe_divide(inBOff, 4), safe_divide(outOff, 4),
        ne00, ne01, ne02,
        nb00, nb01, nb02,
        ne10, ne11, ne12,
        nb10, nb11, nb12,
        ne0, ne1,
        r2, r3
    };

    const unsigned ny = unsigned((ne11 + 4 - 1)/4);

    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(__func__)) {
        const uint32_t local_x = ggml_vk_current_device().subgroupSize * 2;
        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});
    } else {
        s_algo = komputeManager()->getAlgorithm(__func__);
        s_algo->setTensors({inA, inB, out});
        s_algo->setWorkgroup({unsigned(ne01), ny, unsigned(ne12*ne13)});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

static void ggml_vk_mul_mat_mat_f32(kp::Sequence& seq,
                         const std::shared_ptr<kp::Tensor>& inA,
                         const std::shared_ptr<kp::Tensor>& inB,
                         const std::shared_ptr<kp::Tensor>& out,
                         uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
                         int32_t ne00, int32_t ne01, int32_t ne02,
                         uint32_t nb01, uint32_t nb02,
                         int32_t ne11, int32_t ne12,
                         uint32_t nb11, uint32_t nb12,
                         uint32_t nb1, uint32_t nb2) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_mat_f32_comp_spv,
        kp::shader_data::op_mul_mat_mat_f32_comp_spv_len);

    struct PushConstants {
        uint32_t inAOff, inBOff, outOff;
        int32_t ne00, ne01, ne02, ne11, ne12;
        uint32_t nb01, nb02;
        uint32_t nb11, nb12;
        uint32_t nb1, nb2;
    } pushConsts {
        safe_divide(inAOff, 4), safe_divide(inBOff, 4), safe_divide(outOff, 4),
        ne00, ne01, ne02, ne11, ne12,
        nb01, nb02, nb11, nb12,
        nb1, nb2
    };

    const uint32_t local_x = ggml_vk_current_device().subgroupSize;
    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(__func__)) {
        s_algo = komputeManager()->algorithm<uint32_t, PushConstants>(__func__, s_kompute_context->pool.get(),
        {inA, inB, out}, spirv,
        {unsigned(ne01),
         unsigned(ne11),
         unsigned(std::max(ne12, ne02))
         },
        {local_x},
        {pushConsts});
    } else {
        s_algo = komputeManager()->getAlgorithm(__func__);
        s_algo->setTensors({inA, inB, out});
        s_algo->setWorkgroup({unsigned(ne01),
                              unsigned(ne11),
                              unsigned(std::max(ne12, ne02)),
                              });
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

static void ggml_vk_mul_mat_impl(
    const std::vector<uint32_t>& spirv, const char * suffix, uint32_t block_size, kp::Sequence& seq,
    const std::shared_ptr<kp::Tensor>& inA,
    const std::shared_ptr<kp::Tensor>& inB,
    const std::shared_ptr<kp::Tensor>& out,
    uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
    int32_t ne00, int32_t ne01, int32_t ne02,
    int32_t ne10, int32_t ne11, int32_t ne12, int32_t ne13,
    int32_t ne0, int32_t ne1,
    uint32_t r2, uint32_t r3
) {
    struct PushConstants {
        uint32_t inAOff, inBOff, outOff;
        int32_t ne00, ne01, ne02;
        int32_t ne10, ne12;
        int32_t ne0, ne1;
        uint32_t r2, r3;
    } pushConsts {
        safe_divide(inAOff, block_size), safe_divide(inBOff, 4), safe_divide(outOff, 4),
        ne00, ne01, ne02,
        ne10, ne12,
        ne0, ne1,
        r2, r3
    };

    auto name = std::string(__func__) + "_" + suffix;
    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(name)) {
        const uint32_t local_x = ggml_vk_current_device().subgroupSize * 2;
        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});
    } else {
        s_algo = komputeManager()->getAlgorithm(name);
        s_algo->setTensors({inA, inB, out});
        s_algo->setWorkgroup({unsigned((ne01 + 7)/8), unsigned(ne11), unsigned(ne12*ne13)});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

template <typename... Args>
static void ggml_vk_mul_mat_q4_0(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q4_0_comp_spv,
        kp::shader_data::op_mul_mat_q4_0_comp_spv_len);

    ggml_vk_mul_mat_impl(spirv, "q4_0", 1/*We access blocks unaligned*/, std::forward<Args>(args)...);
}

template <typename... Args>
static void ggml_vk_mul_mat_q4_1(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q4_1_comp_spv,
        kp::shader_data::op_mul_mat_q4_1_comp_spv_len);

    ggml_vk_mul_mat_impl(spirv, "q4_1", 1/*We access blocks unaligned*/, std::forward<Args>(args)...);
}

template <typename... Args>
static void ggml_vk_mul_mat_q8_0(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q8_0_comp_spv,
        kp::shader_data::op_mul_mat_q8_0_comp_spv_len);

    ggml_vk_mul_mat_impl(spirv, "q8_0", 1/*We access blocks unaligned*/, std::forward<Args>(args)...);
}

static void ggml_vk_mul_mat_q6_k(
    kp::Sequence& seq,
    const std::shared_ptr<kp::Tensor>& inA,
    const std::shared_ptr<kp::Tensor>& inB,
    const std::shared_ptr<kp::Tensor>& out,
    uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
    int32_t ne00, int32_t ne10, int32_t ne0, int32_t ne1,
    int32_t ne01, int32_t ne11, int32_t ne12, int32_t ne02
) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q6_k_comp_spv,
        kp::shader_data::op_mul_mat_q6_k_comp_spv_len);

    struct PushConstants {
        uint32_t inAOff, inBOff, outOff;
        int32_t ne00, ne10, ne0, ne1, ne01, gqa;
    } pushConsts {
        inAOff, safe_divide(inBOff, 4), safe_divide(outOff, 4),
        ne00, ne10, ne0, ne1, ne01, ne12/ne02
    };

    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(__func__)) {
        const uint32_t local_x = ggml_vk_current_device().subgroupSize * 2;
        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)}, {local_x}, {pushConsts});
    } else {
        s_algo = komputeManager()->getAlgorithm(__func__);
        s_algo->setTensors({inA, inB, out});
        s_algo->setWorkgroup({unsigned((ne01 + 1)/2), unsigned(ne11), unsigned(ne12)});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

static void ggml_vk_get_rows(
    const std::vector<uint32_t>& spirv,
    const char * suffix,
    unsigned element_size, unsigned qk,
    kp::Sequence& seq,
    const std::shared_ptr<kp::Tensor>& inA,
    const std::shared_ptr<kp::Tensor>& inB,
    const std::shared_ptr<kp::Tensor>& out,
    uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
    int32_t ne00, int32_t nb01, int32_t nb1,
    uint32_t size
) {
    GGML_ASSERT(nb01%element_size == 0);
    GGML_ASSERT(nb1%sizeof(float) == 0);
    if (qk) GGML_ASSERT(ne00%qk == 0);

    struct PushConstants {
        uint32_t inAOff, inBOff, outOff;
        int32_t ne00, nb01, nb1;
    } pushConsts {
        safe_divide(inAOff, element_size), safe_divide(inBOff, 4), safe_divide(outOff, 4),
        ne00, nb01, nb1
    };

    auto name = std::string(__func__) + "_" + suffix;
    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(name)) {
        s_algo = komputeManager()->algorithm<float, PushConstants>(name, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {size}, {}, {pushConsts});
    } else {
        s_algo = komputeManager()->getAlgorithm(name);
        s_algo->setTensors({inA, inB, out});
        s_algo->setWorkgroup({size});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

template <typename... Args>
static void ggml_vk_get_rows_f16(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_getrows_f16_comp_spv,
        kp::shader_data::op_getrows_f16_comp_spv_len);

    ggml_vk_get_rows(spirv, "f16", sizeof(half), 0, std::forward<Args>(args)...);
}

template <typename... Args>
static void ggml_vk_get_rows_q4_0(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_getrows_q4_0_comp_spv,
        kp::shader_data::op_getrows_q4_0_comp_spv_len);

    ggml_vk_get_rows(spirv, "q4_0", 1/*We access blocks unaligned*/, QK4_0, std::forward<Args>(args)...);
}

template <typename... Args>
static void ggml_vk_get_rows_q4_1(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_getrows_q4_1_comp_spv,
        kp::shader_data::op_getrows_q4_1_comp_spv_len);

    ggml_vk_get_rows(spirv, "q4_1", 1/*We access blocks unaligned*/, QK4_1, std::forward<Args>(args)...);
}

template <typename... Args>
static void ggml_vk_get_rows_q6_k(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_getrows_q6_k_comp_spv,
        kp::shader_data::op_getrows_q6_k_comp_spv_len);
    ggml_vk_get_rows(spirv, "q6_k", 1/*We access blocks unaligned*/, QK_NL, std::forward<Args>(args)...);
}

static void ggml_vk_rope(
    kp::Sequence& seq,
    const std::shared_ptr<kp::Tensor>& inA,
    const std::shared_ptr<kp::Tensor>& inB,
    const std::shared_ptr<kp::Tensor>& out,
    uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
    ggml_type src0t, int32_t n_dims, int32_t mode, int32_t n_orig_ctx,
    float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow,
    int32_t ne01, int32_t ne02, int32_t ne03,
    uint32_t nb00, uint32_t nb01, uint32_t nb02, uint32_t nb03,
    int32_t ne0,
    uint32_t nb0, uint32_t nb1, uint32_t nb2, uint32_t nb3
) {
    GGML_ASSERT(src0t == GGML_TYPE_F16 || src0t == GGML_TYPE_F32);

    static const auto spirv_f16 = getSpirvShader(
        kp::shader_data::op_rope_f16_comp_spv, kp::shader_data::op_rope_f16_comp_spv_len
    );
    static const auto spirv_f32 = getSpirvShader(
        kp::shader_data::op_rope_f32_comp_spv, kp::shader_data::op_rope_f32_comp_spv_len
    );

    int type_size = src0t == GGML_TYPE_F16 ? 2 : 4;

    GGML_ASSERT(nb03 % type_size == 0);
    GGML_ASSERT(nb02 % type_size == 0);
    GGML_ASSERT(nb01 % type_size == 0);
    GGML_ASSERT(nb00 % type_size == 0);
    GGML_ASSERT(nb3  % type_size == 0);
    GGML_ASSERT(nb2  % type_size == 0);
    GGML_ASSERT(nb1  % type_size == 0);
    GGML_ASSERT(nb0  % type_size == 0);

    struct PushConstants {
        uint32_t inAOff, inBOff, outOff;
        int32_t n_dims, mode, n_orig_ctx;
        float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow;
        uint32_t nb00, nb01, nb02, nb03;
        int32_t ne0;
        uint32_t nb0, nb1, nb2, nb3;
    } pushConsts {
        safe_divide(inAOff, type_size), safe_divide(inBOff, 4), safe_divide(outOff, type_size),
        n_dims, mode, n_orig_ctx,
        freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow,
        nb00, nb01, nb02, nb03,
        ne0,
        nb0, nb1, nb2, nb3
    };

    auto name = std::string(__func__) + (src0t == GGML_TYPE_F16 ? "_f16" : "_f32");
    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(name)) {
        s_algo = komputeManager()->algorithm<float, PushConstants>(
            name, s_kompute_context->pool.get(), {inA, inB, out},
            src0t == GGML_TYPE_F16 ? spirv_f16 : spirv_f32,
            {unsigned(ne01), unsigned(ne02), unsigned(ne03)}, {}, {pushConsts}
        );
    } else {
        s_algo = komputeManager()->getAlgorithm(name);
        s_algo->setTensors({inA, inB, out});
        s_algo->setWorkgroup({unsigned(ne01), unsigned(ne02), unsigned(ne03)});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

static void ggml_vk_cpy(
    const std::vector<uint32_t>& spirv,
    uint32_t in_element_size, uint32_t out_element_size,
    kp::Sequence& seq,
    const std::shared_ptr<kp::Tensor>& in,
    const std::shared_ptr<kp::Tensor>& out,
    uint32_t inOff, uint32_t outOff,
    int32_t ne00, int32_t ne01, int32_t ne02, int32_t ne03,
    uint32_t nb00, uint32_t nb01, uint32_t nb02, uint32_t nb03,
    int32_t ne0, int32_t ne1, int32_t ne2,
    uint32_t nb0, uint32_t nb1, uint32_t nb2, uint32_t nb3
) {
    struct PushConstants {
        uint32_t inOff, outOff;
        int32_t ne00, ne01, ne02;
        uint32_t nb00, nb01, nb02, nb03;
        int32_t ne0, ne1, ne2;
        uint32_t nb0, nb1, nb2, nb3;
    } pushConsts {
        safe_divide(inOff, in_element_size), safe_divide(outOff, out_element_size),
        ne00, ne01, ne02,
        nb00, nb01, nb02, nb03,
        ne0, ne1, ne2,
        nb0, nb1, nb2, nb3
    };

    std::string name = std::string(__func__)
                       + "_i_" + std::to_string(in_element_size)
                       + "_o_" + std::to_string(out_element_size);
    std::shared_ptr<kp::Algorithm> s_algo = nullptr;
    if (!komputeManager()->hasAlgorithm(name))
        s_algo = komputeManager()->algorithm<float, PushConstants>(name, s_kompute_context->pool.get(), {in, out}, spirv, {unsigned(ne01), unsigned(ne02), unsigned(ne03)}, {}, {pushConsts});
    else {
        s_algo = komputeManager()->getAlgorithm(name);
        s_algo->setTensors({in, out});
        s_algo->setWorkgroup({unsigned(ne01), unsigned(ne02), unsigned(ne03)});
        s_algo->setPushConstants<PushConstants>({pushConsts});
        s_algo->updateDescriptors(s_kompute_context->pool.get());
    }
    seq.record<kp::OpAlgoDispatch>(s_algo);
}

template <typename... Args>
static void ggml_vk_cpy_f32_f16(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_cpy_f32_f16_comp_spv,
        kp::shader_data::op_cpy_f32_f16_comp_spv_len);
    ggml_vk_cpy(spirv, 4, 2, std::forward<Args>(args)...);
}

template <typename... Args>
static void ggml_vk_cpy_f32_f32(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_cpy_f32_f32_comp_spv,
        kp::shader_data::op_cpy_f32_f32_comp_spv_len);
    ggml_vk_cpy(spirv, 4, 4, std::forward<Args>(args)...);
}

template <typename... Args>
static void ggml_vk_cpy_f16_f16(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_cpy_f16_f16_comp_spv,
        kp::shader_data::op_cpy_f16_f16_comp_spv_len);
    ggml_vk_cpy(spirv, 2, 2, std::forward<Args>(args)...);
}

template <typename... Args>
static void ggml_vk_cpy_f16_f32(Args&&... args) {
    const static auto spirv = getSpirvShader(kp::shader_data::op_cpy_f16_f32_comp_spv,
        kp::shader_data::op_cpy_f16_f32_comp_spv_len);
    ggml_vk_cpy(spirv, 2, 4, std::forward<Args>(args)...);
}

static bool ggml_vk_supports_op(const struct ggml_tensor * op) {
    switch (op->type) {
        case GGML_TYPE_F16:
        case GGML_TYPE_F32:
        case GGML_TYPE_Q4_0:
        case GGML_TYPE_Q4_1:
            break;
        default:
            return false;
    }

    switch (op->op) {
        case GGML_OP_UNARY:
            switch (ggml_get_unary_op(op)) {
                case GGML_UNARY_OP_RELU:
                case GGML_UNARY_OP_GELU:
                case GGML_UNARY_OP_SILU:
                    return true;
                default:
                    ;
            }
            break;
        case GGML_OP_NONE:
        case GGML_OP_RESHAPE:
        case GGML_OP_VIEW:
        case GGML_OP_TRANSPOSE:
        case GGML_OP_PERMUTE:
        case GGML_OP_ADD:
        case GGML_OP_MUL:
        case GGML_OP_SCALE:
        case GGML_OP_SOFT_MAX:
        case GGML_OP_RMS_NORM:
        case GGML_OP_NORM:
        case GGML_OP_ROPE:
            return true;
        case GGML_OP_DUP:
        case GGML_OP_CPY:
        case GGML_OP_CONT:
            switch (op->src[0]->type) {
                case GGML_TYPE_F32:
                case GGML_TYPE_F16:
                    break;
                default:
                    return false;
            }
            switch (op->type) {
                case GGML_TYPE_F32:
                case GGML_TYPE_F16:
                    break;
                default:
                    return false;
            }
            return true;
        case GGML_OP_DIAG_MASK_INF:
            return op->ne[3] == 1;
        case GGML_OP_GET_ROWS:
            switch (op->src[0]->type) {
                case GGML_TYPE_F16:
                case GGML_TYPE_Q4_0:
                case GGML_TYPE_Q4_1:
                case GGML_TYPE_Q6_K:
                    return op->ne[2] == 1 && op->ne[3] == 1;
                default:
                    ;
            }
            return false;
        case GGML_OP_MUL_MAT:
            if (op->src[1]->type != GGML_TYPE_F32 || ggml_is_transposed(op->src[0]) || ggml_is_transposed(op->src[1]))
                return false;

            switch (op->src[0]->type) {
                case GGML_TYPE_F32:
                case GGML_TYPE_Q6_K:
                    return op->ne[3] == 1;
                case GGML_TYPE_F16:
                case GGML_TYPE_Q8_0:
                case GGML_TYPE_Q4_0:
                case GGML_TYPE_Q4_1:
                    return true;
                default:
                    ;
            }
        default:
            ;
    }
    return false;
}

static void ggml_vk_graph_compute(struct ggml_kompute_context * ctx, struct ggml_cgraph * gf) {
    const int n_seq = 8;

    // FIXME: Figure out if we can somehow optimize the size of the pool... right now we're setting
    // it to the size of the graph, but I think it can be made smaller?
    ggml_vk_allocate_descriptor_pool(ctx, gf->n_nodes);

    std::vector<std::shared_ptr<kp::Sequence>> sequences(n_seq);

    for (auto& sequence : sequences) {
        sequence = komputeManager()->sequence();
    }
    for (int seq_idx = 0; seq_idx < n_seq; ++seq_idx) {
        const int n_nodes_per_seq = (gf->n_nodes + n_seq - 1) / n_seq;

        auto& seq = *sequences[seq_idx];

        const int node_start = (seq_idx + 0) * n_nodes_per_seq;
        const int node_end   = std::min((seq_idx == n_seq - 1) ? gf->n_nodes : (seq_idx + 1) * n_nodes_per_seq, gf->n_nodes);

        bool any_commands_recorded = false;

        for (int i = node_start; i < node_end; ++i) {
            struct ggml_tensor * src0 = gf->nodes[i]->src[0];
            struct ggml_tensor * src1 = gf->nodes[i]->src[1];
            struct ggml_tensor * dst = gf->nodes[i];
            GGML_ASSERT(dst->data != nullptr);

            switch (dst->op) {
                case GGML_OP_NONE:
                case GGML_OP_RESHAPE:
                case GGML_OP_VIEW:
                case GGML_OP_TRANSPOSE:
                case GGML_OP_PERMUTE:
                    continue; // noop -> next node
                default:
                    break;
            }

            any_commands_recorded = true;

            if (!ggml_vk_supports_op(dst)) {
                 fprintf(stderr, "%s: error: unsupported op '%s'\n", __func__, ggml_op_desc(dst));
                 GGML_ASSERT(!"unsupported op");
            }

            const int32_t ne00 = src0 ? src0->ne[0] : 0;
            const int32_t ne01 = src0 ? src0->ne[1] : 0;
            const int32_t ne02 = src0 ? src0->ne[2] : 0;
            const int32_t ne03 = src0 ? src0->ne[3] : 0;

            const uint32_t nb00 = src0 ? src0->nb[0] : 0;
            const uint32_t nb01 = src0 ? src0->nb[1] : 0;
            const uint32_t nb02 = src0 ? src0->nb[2] : 0;
            const uint32_t nb03 = src0 ? src0->nb[3] : 0;

            const int32_t ne10 = src1 ? src1->ne[0] : 0;
            const int32_t ne11 = src1 ? src1->ne[1] : 0;
            const int32_t ne12 = src1 ? src1->ne[2] : 0;
            const int32_t ne13 = src1 ? src1->ne[3] : 0;

            const uint32_t nb10 = src1 ? src1->nb[0] : 0;
            const uint32_t nb11 = src1 ? src1->nb[1] : 0;
            const uint32_t nb12 = src1 ? src1->nb[2] : 0;
            const uint32_t nb13 = src1 ? src1->nb[3] : 0;

            const int32_t ne0 = dst ? dst->ne[0] : 0;
            const int32_t ne1 = dst ? dst->ne[1] : 0;
            const int32_t ne2 = dst ? dst->ne[2] : 0;
//            const int32_t ne3 = dst ? dst->ne[3] : 0;

            const uint32_t nb0 = dst ? dst->nb[0] : 0;
            const uint32_t nb1 = dst ? dst->nb[1] : 0;
            const uint32_t nb2 = dst ? dst->nb[2] : 0;
            const uint32_t nb3 = dst ? dst->nb[3] : 0;

            const enum ggml_type src0t = src0 ? src0->type : GGML_TYPE_COUNT;
            const enum ggml_type src1t = src1 ? src1->type : GGML_TYPE_COUNT;
            const enum ggml_type dstt = dst ? dst->type : GGML_TYPE_COUNT;

            uint32_t off_src0 = 0;
            uint32_t off_src1 = 0;
            uint32_t off_dst  = 0;
            const std::shared_ptr<kp::Tensor>& id_src0 = src0 ? ggml_vk_get_tensor_aligned(src0, &off_src0) : nullTensor;
            const std::shared_ptr<kp::Tensor>& id_src1 = src1 ? ggml_vk_get_tensor_aligned(src1, &off_src1) : nullTensor;
            const std::shared_ptr<kp::Tensor>& id_dst  = dst  ? ggml_vk_get_tensor_aligned(dst,  &off_dst)  : nullTensor;

            switch (dst->op) {
                case GGML_OP_ADD:
                    {
                        if (ggml_nelements(src1) == ne10 && ggml_is_contiguous(src1) && ne00 % 4 == 0 && ne10 % 4 == 0) {
                            // src1 is a row
                            ggml_vk_addrow(seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst, ggml_nelements(dst)/4, ne00);
                        } else {
                            ggml_vk_add(
                                seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst,
                                ne00, ne01, ne02, ne03,
                                nb00, nb01, nb02, nb03,
                                ne10, ne11, ne12, ne13,
                                nb10, nb11, nb12, nb13,
                                ne0,
                                nb0, nb1, nb2, nb3
                            );
                        }
                    } break;
                case GGML_OP_MUL:
                    {
                        ggml_vk_mul(
                            seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst,
                            ne00, ne01, ne02, ne03,
                            nb00, nb01, nb02, nb03,
                            ne10, ne11, ne12, ne13,
                            nb10, nb11, nb12, nb13,
                            ne0,
                            nb0, nb1, nb2, nb3
                        );
                    } break;
                case GGML_OP_SCALE:
                    {
                        float scale; memcpy(&scale, dst->op_params, sizeof(float));

                        ggml_vk_scale(seq, id_src0, id_dst, off_src0, off_dst, ggml_nelements(dst), scale);
                    } break;
                case GGML_OP_UNARY:
                    {
                        int64_t n = ggml_nelements(dst);
                        GGML_ASSERT(n % 4 == 0);
                        switch (ggml_get_unary_op(gf->nodes[i])) {
                            case GGML_UNARY_OP_SILU:
                                {
                                    ggml_vk_silu(seq, id_src0, id_dst, off_src0, off_dst, n/4);
                                } break;
                            case GGML_UNARY_OP_RELU:
                                {
                                    ggml_vk_relu(seq, id_src0, id_dst, off_src0, off_dst, n/4);
                                } break;
                            case GGML_UNARY_OP_GELU:
                                {
                                    GGML_ASSERT(n % 8 == 0);
                                    ggml_vk_gelu(seq, id_src0, id_dst, off_src0, off_dst, n/8);
                                } break;
                            default:
                                {
                                    fprintf(stderr, "%s: node %3d, op = %8s not implemented\n", __func__, i, ggml_op_name(dst->op));
                                    GGML_ASSERT(false);
                                }
                        }
                    } break;
                case GGML_OP_SOFT_MAX:
                    {
                        float scale;
                        memcpy(&scale, dst->op_params, sizeof(float));
                        ggml_vk_soft_max(seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst, ne00, ne01, ne02, ne03, scale);
                    } break;
                case GGML_OP_DIAG_MASK_INF:
                    {
                        const int n_past = ((int32_t *)(dst->op_params))[0];
                        ggml_vk_diag_mask_inf(seq, id_src0, id_dst, off_src0, off_dst, n_past, ne00, ne01, ne02);
                    } break;
                case GGML_OP_NORM:
                    {
                        float eps;
                        memcpy(&eps, dst->op_params, sizeof(float));
                        ggml_vk_norm(seq, id_src0, id_dst, off_src0, off_dst, ne00, nb01, ggml_nrows(src0), eps);
                    } break;
                case GGML_OP_RMS_NORM:
                    {
                        GGML_ASSERT(ne00 % 4 == 0);

                        float eps;
                        memcpy(&eps, dst->op_params, sizeof(float));
                        ggml_vk_rms_norm(seq, id_src0, id_dst, off_src0, off_dst, ne00, nb01, ggml_nrows(src0), eps);
                    } break;
                case GGML_OP_MUL_MAT:
                    {
                        GGML_ASSERT(ne00 == ne10);

                        // TODO: assert that dim2 and dim3 are contiguous
                        GGML_ASSERT(ne12 % ne02 == 0);
                        GGML_ASSERT(ne13 % ne03 == 0);

                        const uint32_t r2 = ne12/ne02;
                        const uint32_t r3 = ne13/ne03;

                        if (src1t != GGML_TYPE_F32) {
                            fprintf(stderr, "%s: %s: Unsupported src1 type: %u/%u\n", __func__, ggml_op_name(dst->op), src0t, src1t);
                            goto not_implemented;
                        }

                        if (ggml_is_transposed(src0) ||
                            ggml_is_transposed(src1)) {
                            fprintf(stderr, "%s: %s: matmul on tranposed tensor not supported: %u/%u\n", __func__, ggml_op_name(dst->op), src0t, src1t);
                            goto not_implemented;
                        }

                        switch (src0t) {
                            case GGML_TYPE_F32:
                                ggml_vk_mul_mat_mat_f32(
                                    seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst,
                                    ne00, ne01, ne02, nb01, nb02, ne11, ne12, nb11, nb12, nb1, nb2
                                );
                                break;
                            case GGML_TYPE_F16:
                                ggml_vk_mul_mat_f16(
                                    seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst,
                                    ne00, ne01, ne02, nb00, nb01, nb02, ne10, ne11, ne12, ne13, nb10, nb11, nb12,
                                    ne0, ne1, r2, r3
                                );
                                break;
                            case GGML_TYPE_Q8_0:
                                ggml_vk_mul_mat_q8_0(
                                    seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst,
                                    ne00, ne01, ne02, ne10, ne11, ne12, ne13, ne0, ne1, r2, r3
                                );
                                break;
                            case GGML_TYPE_Q4_0:
                                ggml_vk_mul_mat_q4_0(
                                    seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst,
                                    ne00, ne01, ne02, ne10, ne11, ne12, ne13, ne0, ne1, r2, r3
                                );
                                break;
                            case GGML_TYPE_Q4_1:
                                ggml_vk_mul_mat_q4_1(
                                    seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst,
                                    ne00, ne01, ne02, ne10, ne11, ne12, ne13, ne0, ne1, r2, r3
                                );
                                break;
                            case GGML_TYPE_Q6_K:
                                ggml_vk_mul_mat_q6_k(
                                    seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst,
                                    ne00, ne10, ne0, ne1, ne01, ne11, ne12, ne02
                                );
                                break;
                            default: {
                                fprintf(stderr, "%s: %s: Unsupported quantization: %u/%u\n", __func__, ggml_op_name(dst->op), src0t, src1t);
                                goto not_implemented;
                            }
                        }

                    } break;
                case GGML_OP_GET_ROWS:
                    {
                        if (src0t == GGML_TYPE_F16) {
                            ggml_vk_get_rows_f16(seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst, ne00, nb01, nb1, ggml_nelements(src1));
                        } else if (src0t == GGML_TYPE_Q4_0) {
                            ggml_vk_get_rows_q4_0(seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst, ne00, nb01, nb1, ggml_nelements(src1));
                        } else if (src0t == GGML_TYPE_Q4_1) {
                            ggml_vk_get_rows_q4_1(seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst, ne00, nb01, nb1, ggml_nelements(src1));
                        } else if (src0t == GGML_TYPE_Q6_K) {
                            ggml_vk_get_rows_q6_k(seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst, ne00, nb01, nb1, ggml_nelements(src1));
                        } else {
                            fprintf(stderr, "%s: %s: Unsupported quantization: %u\n", __func__, ggml_op_name(dst->op), src0t);
                            goto not_implemented;
                        }
                    } break;
                case GGML_OP_ROPE:
                    {
                        GGML_ASSERT(ne10 == ne02);
                        GGML_ASSERT(src0t == dstt);
                        // const int n_past = ((int32_t *) dst->op_params)[0];
                        const int n_dims     = ((int32_t *) dst->op_params)[1];
                        const int mode       = ((int32_t *) dst->op_params)[2];
                        // skip 3, n_ctx used in GLM RoPE, unimplemented in Vulkan
                        const int n_orig_ctx = ((int32_t *) dst->op_params)[4];

                        float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow;
                        memcpy(&freq_base,   (int32_t *) dst->op_params +  5, sizeof(float));
                        memcpy(&freq_scale,  (int32_t *) dst->op_params +  6, sizeof(float));
                        memcpy(&ext_factor,  (int32_t *) dst->op_params +  7, sizeof(float));
                        memcpy(&attn_factor, (int32_t *) dst->op_params +  8, sizeof(float));
                        memcpy(&beta_fast,   (int32_t *) dst->op_params +  9, sizeof(float));
                        memcpy(&beta_slow,   (int32_t *) dst->op_params + 10, sizeof(float));
                        ggml_vk_rope(
                            seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst, src0t, n_dims, mode, n_orig_ctx,
                            freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow,
                            ne01, ne02, ne03, nb00, nb01, nb02, nb03, ne0, nb0, nb1, nb2, nb3
                        );
                    } break;
                case GGML_OP_DUP:
                case GGML_OP_CPY:
                case GGML_OP_CONT:
                    {
                        switch (src0t) {
                            case GGML_TYPE_F32:
                                {
                                    switch (dstt) {
                                        case GGML_TYPE_F16: ggml_vk_cpy_f32_f16(seq, id_src0, id_dst, off_src0, off_dst, ne00, ne01, ne02, ne03, nb00, nb01, nb02, nb03, ne0, ne1, ne2, nb0, nb1, nb2, nb3); break;
                                        case GGML_TYPE_F32: ggml_vk_cpy_f32_f32(seq, id_src0, id_dst, off_src0, off_dst, ne00, ne01, ne02, ne03, nb00, nb01, nb02, nb03, ne0, ne1, ne2, nb0, nb1, nb2, nb3); break;
                                        default: goto not_implemented;
                                    }
                                } break;
                            case GGML_TYPE_F16:
                                {
                                    switch (dstt) {
                                        case GGML_TYPE_F16: ggml_vk_cpy_f16_f16(seq, id_src0, id_dst, off_src0, off_dst, ne00, ne01, ne02, ne03, nb00, nb01, nb02, nb03, ne0, ne1, ne2, nb0, nb1, nb2, nb3); break;
                                        case GGML_TYPE_F32: ggml_vk_cpy_f16_f32(seq, id_src0, id_dst, off_src0, off_dst, ne00, ne01, ne02, ne03, nb00, nb01, nb02, nb03, ne0, ne1, ne2, nb0, nb1, nb2, nb3); break;
                                    default: goto not_implemented;
                                } break;
                            default: goto not_implemented;
                            }
                        }
                    } break;
                default: goto not_implemented;
            }
            continue;
            not_implemented: {}
            fprintf(stderr, "%s: node %3d, op = %8s not implemented\n", __func__, i, ggml_op_name(dst->op));
            //GGML_ASSERT(false);
        }

        // Evaluate sequence
        if (any_commands_recorded) {
            seq.evalAsync();
        }
    }

    // Wait for all sequences to finish
    for (auto& sequence : sequences) {
        if (sequence->isRunning())
            sequence->evalAwait();
    }

    ggml_vk_free_descriptor_pool(ctx);
}

template<>
kp::Tensor::TensorDataTypes
kp::TensorT<half>::dataType()
{
    return TensorDataTypes::eFloat;
}

template<>
kp::Tensor::TensorDataTypes
kp::TensorT<uint8_t>::dataType()
{
    return TensorDataTypes::eUnsignedInt;
}

////////////////////////////////////////////////////////////////////////////////

// backend interface

struct ggml_backend_kompute_buffer_type_context {
    int         device;
    int         device_ref = 0;
    uint64_t    buffer_alignment;
    uint64_t    max_alloc;
    std::string name;

    ggml_backend_kompute_buffer_type_context(int device, uint64_t buffer_alignment, uint64_t max_alloc)
        : device(device), buffer_alignment(buffer_alignment), max_alloc(max_alloc), name(ggml_kompute_format_name(device)) {}
};

static void ggml_backend_kompute_device_ref(ggml_backend_buffer_type_t buft) {
    auto * ctx = static_cast<ggml_backend_kompute_buffer_type_context *>(buft->context);

    if (!ctx->device_ref) {
        komputeManager()->initializeDevice(
            ctx->device, {}, {"VK_KHR_8bit_storage", "VK_KHR_16bit_storage", "VK_KHR_shader_non_semantic_info"}
        );
    }

    assert(ggml_vk_has_device());
    ctx->device_ref++;
}

static void ggml_backend_kompute_device_unref(ggml_backend_buffer_type_t buft) {
    auto * ctx = static_cast<ggml_backend_kompute_buffer_type_context *>(buft->context);

    assert(ctx->device_ref > 0);

    ctx->device_ref--;

    if (!ctx->device_ref) {
        komputeManager.destroy();
    }
}

static const char * ggml_backend_kompute_buffer_get_name(ggml_backend_buffer_t buffer) {
    auto * ctx = static_cast<ggml_backend_kompute_buffer_type_context *>(buffer->buft->context);
    return ctx->name.c_str();
}

static void ggml_backend_kompute_buffer_free_buffer(ggml_backend_buffer_t buffer) {
    auto * memory = (ggml_vk_memory *)buffer->context;
    ggml_vk_free_memory(*memory);
    delete memory;
    ggml_backend_kompute_device_unref(buffer->buft);
}

static void * ggml_backend_kompute_buffer_get_base(ggml_backend_buffer_t buffer) {
    return ((ggml_vk_memory *)buffer->context)->data;
}

static void ggml_backend_kompute_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
    GGML_UNUSED(buffer);

    const auto res = ggml_vk_get_tensor_slice(tensor, offset, size);
    GGML_ASSERT(res);

    memcpy((char *)tensor->data + offset, data, size);

    komputeManager()->sequence()->eval<kp::OpTensorSyncDevice>({res});
}

static void ggml_backend_kompute_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
    GGML_UNUSED(buffer);

    const auto res = ggml_vk_get_tensor_slice(tensor, offset, size);
    GGML_ASSERT(res);

    komputeManager()->sequence()->eval<kp::OpTensorSyncLocal>({res});

    memcpy(data, (const char *)tensor->data + offset, size);
}

static void ggml_backend_kompute_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
    auto * memory = (ggml_vk_memory *)buffer->context;
    memset(memory->data, value, buffer->size);

    if (memory->stagingBuffer)
        komputeManager()->sequence()->eval<kp::OpBufferSyncDevice>(memory->primaryBuffer, memory->stagingBuffer, memory->size);
}

static ggml_backend_buffer_i ggml_backend_kompute_buffer_i = {
    /* .get_name        = */ ggml_backend_kompute_buffer_get_name,
    /* .free_buffer     = */ ggml_backend_kompute_buffer_free_buffer,
    /* .get_base        = */ ggml_backend_kompute_buffer_get_base,
    /* .init_tensor     = */ NULL,
    /* .set_tensor      = */ ggml_backend_kompute_buffer_set_tensor,
    /* .get_tensor      = */ ggml_backend_kompute_buffer_get_tensor,
    /* .cpy_tensor      = */ NULL,
    /* .clear           = */ ggml_backend_kompute_buffer_clear,
    /* .reset           = */ NULL,
};

// default buffer type

static const char * ggml_backend_kompute_buffer_type_get_name(ggml_backend_buffer_type_t buft) {
    auto * ctx = static_cast<ggml_backend_kompute_buffer_type_context *>(buft->context);
    return ctx->name.c_str();
}

static ggml_backend_buffer_t ggml_backend_kompute_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
    ggml_backend_kompute_device_ref(buft);
    auto * ctx = new ggml_vk_memory(ggml_vk_allocate(size));
    if (!ctx->primaryMemory) {
        delete ctx;
        return nullptr; // allocation failed
    }
    return ggml_backend_buffer_init(buft, ggml_backend_kompute_buffer_i, ctx, size);
}

static size_t ggml_backend_kompute_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
    auto * ctx = static_cast<ggml_backend_kompute_buffer_type_context *>(buft->context);
    return ctx->buffer_alignment;
}

static size_t ggml_backend_vk_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) {
    auto * ctx = static_cast<ggml_backend_kompute_buffer_type_context *>(buft->context);
    return ctx->max_alloc;
}

static bool ggml_backend_kompute_buffer_type_supports_backend(ggml_backend_buffer_type_t buft, ggml_backend_t backend) {
    GGML_UNUSED(buft);
    return ggml_backend_is_kompute(backend);
}

static ggml_backend_buffer_type_i ggml_backend_kompute_buffer_type_interface = {
    /* .get_name         = */ ggml_backend_kompute_buffer_type_get_name,
    /* .alloc_buffer     = */ ggml_backend_kompute_buffer_type_alloc_buffer,
    /* .get_alignment    = */ ggml_backend_kompute_buffer_type_get_alignment,
    /* .get_max_size     = */ ggml_backend_vk_buffer_type_get_max_size,
    /* .get_alloc_size   = */ NULL, // defaults to ggml_nbytes
    /* .supports_backend = */ ggml_backend_kompute_buffer_type_supports_backend,
    /* .is_host          = */ NULL,
};

ggml_backend_buffer_type_t ggml_backend_kompute_buffer_type(int device) {
    static std::vector<ggml_backend_buffer_type> bufts = []() {
        std::vector<ggml_backend_buffer_type> vec;
        auto devices = ggml_vk_available_devices_internal(0);
        vec.reserve(devices.size());

        for (const auto & dev : devices) {
            vec.push_back({
                /* .iface   = */ ggml_backend_kompute_buffer_type_interface,
                /* .context = */ new ggml_backend_kompute_buffer_type_context(dev.index, dev.bufferAlignment, dev.maxAlloc)
            });
        }
        return vec;
    }();

    auto it = std::find_if(bufts.begin(), bufts.end(), [device](const ggml_backend_buffer_type & t) {
        return device == static_cast<ggml_backend_kompute_buffer_type_context *>(t.context)->device;
    });
    return it < bufts.end() ? &*it : nullptr;
}

// backend

static const char * ggml_backend_kompute_name(ggml_backend_t backend) {
    auto * ctx = static_cast<ggml_kompute_context *>(backend->context);
    return ctx->name.c_str();
}

static void ggml_backend_kompute_free(ggml_backend_t backend) {
    auto * ctx = static_cast<ggml_kompute_context *>(backend->context);

    assert(ctx == s_kompute_context);
    s_kompute_context = nullptr;
    if (ctx != nullptr) {
        delete ctx;
    }

    delete backend;
}

static ggml_backend_buffer_type_t ggml_backend_kompute_get_default_buffer_type(ggml_backend_t backend) {
    auto * ctx = static_cast<ggml_kompute_context *>(backend->context);
    return ggml_backend_kompute_buffer_type(ctx->device);
}

static ggml_status ggml_backend_kompute_graph_compute(ggml_backend_t backend, struct ggml_cgraph * cgraph) {
    auto * ctx = static_cast<ggml_kompute_context *>(backend->context);
    ggml_vk_graph_compute(ctx, cgraph);
    return GGML_STATUS_SUCCESS;
}

static bool ggml_backend_kompute_supports_op(ggml_backend_t backend, const struct ggml_tensor * op) {
    GGML_UNUSED(backend);
    return ggml_vk_supports_op(op);
}

static struct ggml_backend_i kompute_backend_i = {
    /* .get_name                = */ ggml_backend_kompute_name,
    /* .free                    = */ ggml_backend_kompute_free,
    /* .get_default_buffer_type = */ ggml_backend_kompute_get_default_buffer_type,
    /* .set_tensor_async        = */ NULL,
    /* .get_tensor_async        = */ NULL,
    /* .cpy_tensor_async        = */ NULL,
    /* .synchronize             = */ NULL,
    /* .graph_plan_create       = */ NULL,
    /* .graph_plan_free         = */ NULL,
    /* .graph_plan_compute      = */ NULL,
    /* .graph_compute           = */ ggml_backend_kompute_graph_compute,
    /* .supports_op             = */ ggml_backend_kompute_supports_op,
};

static ggml_guid_t ggml_backend_kompute_guid() {
    static ggml_guid guid = { 0x7b, 0x57, 0xdc, 0xaf, 0xde, 0x12, 0x1d, 0x49, 0xfb, 0x35, 0xfa, 0x9b, 0x18, 0x31, 0x1d, 0xca };
    return &guid;
}

ggml_backend_t ggml_backend_kompute_init(int device) {
    GGML_ASSERT(s_kompute_context == nullptr);
    s_kompute_context = new ggml_kompute_context(device);

    ggml_backend_t kompute_backend = new ggml_backend {
        /* .guid      = */ ggml_backend_kompute_guid(),
        /* .interface = */ kompute_backend_i,
        /* .context   = */ s_kompute_context,
    };

    return kompute_backend;
}

bool ggml_backend_is_kompute(ggml_backend_t backend) {
    return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_kompute_guid());
}

static ggml_backend_t ggml_backend_reg_kompute_init(const char * params, void * user_data) {
    GGML_UNUSED(params);
    return ggml_backend_kompute_init(intptr_t(user_data));
}

extern "C" int ggml_backend_kompute_reg_devices();

int ggml_backend_kompute_reg_devices() {
    auto devices = ggml_vk_available_devices_internal(0);
    for (const auto & device : devices) {
        ggml_backend_register(
            ggml_kompute_format_name(device.index).c_str(),
            ggml_backend_reg_kompute_init,
            ggml_backend_kompute_buffer_type(device.index),
            reinterpret_cast<void *>(intptr_t(device.index))
        );
    }
    return devices.size();
}